Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks
Authors: Felix Dangel
Johannes Müller
Marius Zeinhofer
Venue: NeurIPS 2024 poster (2024)
Keywords: KFAC
PINNs
Gauss-Newton
PDEs
Taylor mode automatic differentiation
Forward Laplacian
Second-order optimization
Higher-order derivatives
Primary area: machine_learning_for_physical_sciences
TL;DR: We derive a KFAC approximation for PINN losses which scales to high-dimensional NNs and PDEs and consistently outperforms first-order methods for training PINNs.

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 popular 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 which allows us to apply a variant of KFAC 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.

Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks
Felix Dangel"
Johannes Müller
Vector Institute
Chair of Mathematics of Information Processing
Toronto
RWTH Aachen University
Canada
Aachen, Germany
fdangee.
mueller@mathc. www..rwwwh..hhcccccc
Marius Zeinhofer"
Seminar for Applied Mathematics, ETH Zürich,
Department of Nuclear Medicine, University Hospital Freiburg de marius aeeinhofrruuiiiiiikkk
Abstract
Physics-informed neural networks (PINNs) are infamous for being hard 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 networks with few thousand parameters due to the high computational cost to evaluate, store, and invert the curvature matrix. We propose Kronecker-factoree 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 it captures contributions from 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 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 -orde methods small problems, scale more favorably to higher-dimensional neural networks and PDEs, and consistently outperform first-order methods and LBFGS.
Introduction
Neural network-basse approaches to numerically solve partial differential equations (PDEs) are growing at unprecedented speed. The idea train network parameters 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 [52], Raissi et al. [50]. PINNs are arguably one of the most popular network- seteork-based approaches 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 [57] and fail to provide satisfactory accuracy when trained with first-order methods, even for simple problems [64, 41]. Recently, second-order methods that use the function space geometry to design preconditioners
"Equal contribution
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

have shown remarkable promise in addressing the training difficulties of PINNs [64, 41, 14, 24, 42]. However, these methods require solving a linear system in the network's nigh-dimensional parameter space at cubic computational iteration cost, which prohibits scaling such approaches. To address this, we build on the idea of Kronecker-factee 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 compu- tation graph of network's s input derivatives as 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 eq. (14)). Thanks to the generality of Taylor-mode and KFAC for weight sharing layers [17], 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 more efficient scheme. For the prominent example of the Laplace operator, this recovers and generalizes the forward Laplacian framework [29] ($3.2 and eq. (9)).
Empirically, we find that our KFAC-based optimizers are competitive with expensive second- order methods on small problems, scale more favorably to nigher-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 [57. 56, 59], different sampling strategies for discretizing the loss [34, 43, 13, 63, 58, 611, and curriculum learning [26, 58]. While LBFGS is known to improve upon first-order optimizers [35], recently, other second-order methods that design meaningful preconditioners that respect the problem's 's geometry have significantly outperformed it [64, 41, 14, 33, 24,7 7, 62]. Müller & Zeinhofer [42] provide unified view on these approaches which greatly improve the accuracy of PINNs, but come with significant per-iteration cost as one needs to solve linear system in the network's ' high-dimens -ichdimmensinal 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 [51, 36] for supervised learning problems and [64, 4, 24, 62) for PINNs. Instead, our KFAC-based approach uses explicit structured curvature representation which can be updated over iterations and inverted more cheaply.
We build on the literature on -rotecker-factored approximate curvature (KFAC), which was initially introduced in Heskes [22], Martens [36] as an approximation of the per-layer Fisher matrix to perform approximate natural gradient descent. Later, KFAC was extended to convolutional [20], recurrent [39], attention [47, 44, 21], and recently general linear layers with weight sharing [17]. These works do not address preconditioners for losses with contributions from differential operators, as is the case for PINN losses. 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 [17].
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 allows 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 post-index-vare varies-faseess convention, iee column-stacking (row index varies first, column index varies second) and denote the corresponding flattening operation by vec. This allows to reduce derivatives of matrix- or tensor-vas tensor-valued objects back to the vector case by flattening a function's 's input and output before differentiation. The Jacobian of a vector-to-vector function a + b(a) has entries [Jab]e,j = 0b./oaj. For a matrix-to- matrix function A + B(A), the Jacobian is J B Jvee A vec B. A useful property of vec is [BT T @ A) vec X for matrices A, X, B which implies (AXB) = B & A.
Sequential neural nets Consider a sequential neural network ug = Sour focc-1) o... o soa:
of depth L E N. It consists of layers fao: - h) 2 (1-1) + z!!) fow (2)02((1-)) (2(1-1)) with
2

---

trainable parameters O(1 E R that transform input 20)) = 2: Rd= into prediction (x(x) =()) E Bh via intermediate representations 2 E R In the context of PINNs, we use networks with scalar outputs (h(z)) 1) and denote the concatenation of all parameters by
) E RD. A common choice is to alternate fully-connected and activation layers. Linear layers map z(l-1) + z()) (() using weight matrix W vec g(0)
p xh(t-1> (bias terms can added as an additional column and by appending a 1 to the input). Activation layers map (!-1) + 211) (2z(1-1) element-wise for a (typically smooth) a: + I.
2.1
Energy Natural Gradients for Physics- Physics-Informee Informed Neural Networks
Let us consider a domain Ω C Rdd and the partial differential equation
Lu= in S2 u= on i.
with right-hand side f, boundary data g and differential operator L, e.g. the negative Laplacian -Lu An Σi 02, u. We parametrize with neural net and train its parameters 0 to minimize the loss
No Σ
Naa (uo(x²) g(xhh))²
L(O) 2No
(Cuu((an) f(xx))² 2N
n=1
=: (0) + Lon
with points {z,, E ? from the domain's interior, and points {жh E a on its boundary."
First-order optimizers like gradient descent and Adam struggle at producing satisfactory solutions when used train PINNs [9]. Instead, function space-inspired second-order methods have lately shown promising results [42]. We focus on energy natural gradient descent (ENGI) [4/]) which- applied to PINN objectives like ())-corrrpponds ENGD mimics Newton's '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 -Nwstn matrix leads to optimal conditioning of problem [14].
the Gauss-Newton method [6, Chapter 6.3].
Natural gradient methods perform parameter updates via a preconditioned gradient descent scheme + 0 (0 ((O)). where G(0)+ denotes the pseudo-inverse of a suitable Gramian matrix G( E RDD and a step size. ENGD for the PINN loss (1) uses the Gramian
Nea (Jouo Jouo(zh)
G( Nn Σ (JoLuo(zn)) JoLuo(In)) Non
(2)
=: (0) + Go)
(2) is the Gauss-Newton matrix of the residual r0) = (ra(0))/VNT)) roa(0) //Noo) T E R NittNon with interior and boundary residuals rn, (0) Lug (a... fxx,) and rant uo (a())) (xx)))).
2.2 Kronecker-factoree Approximate Curvature
We review Kronecker-factored approximate curvature (KFAC) which was introduced by Heskes [22], 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 1/2N En (.)) - yn112 with targets yr. E IR is
N
F(0) =
Σ (J Jgun Σ (Jou Jaue( N ==1 1=1
DDD
(3)
N
where M, uo (xx and it coincides with the classical Gauss-Ner -Nwttn matrix [37]. The established
KFAC approximates (3), While the boundary Gramian G (0) has the same structure as F()), the
interior Gramian ()) does not it involves derivative rather than function evaluations of the net.
Thh second regression loss also include other constraints like measurement data.

---

KFAC tackles the Fisher's per-layer block diagonal, F(0) = diag(F()())
(O) with
(0) 1/v E (J " Join un E R xp! For a fully-connected layer's block, let's examine the term Jø:n up( from Equation (3) for fixed data point. The layer parameters 0
vec W/() enter the computation via z0) w(0) and we have J Jway = ((111)
le.g.
10J. Further, the chain rule gives the decomposition Jw = ww = (1-1) @ Jmm u. Inserting into (0), summing over data points, and using the expectation approximation Σ,, A, ?
B, Re (E, A.) @ (En B,) from Martens & Grosse [38], we obtain the KFAC approximation for linear layers in supervised square loss regression with a network's output.
N
N (Jac) Un) 22
F (0)
1 N
(1-1)
(1-1) n
1
(4)
n=1
=A = xx(t-11)))
=B0 (RA())x())
It is cheap to store and invert by inverting the two Kronecker factors.
Kronecker- Approximate Curvature for PINNs
ENGD's Gramian is a sum of PDE and boundary Gramians, G(Ø) = G)z(0) + Gon We
will approximate each Gramian separately with block diagonal matrix with Kronecker-factore
blocks, G.) diag diig (O).... G. for E {S2, 02) with G! (0) ~ (1) x B!!
For the boundary Gramian Gen (0), we can re-use the established KFAC from Equation (4) as its loss corresponds to regression over the network's s output. The interior Gramian Go however, involves PDE terms in the form of network derivatives and therefore cannor 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, need to make the dependency between the weights and the differential operator Lu explicit. We use Taylor-ma 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 KFAC's recently-pree propposd generalization to linear layers with weight sharing [17].
3.1 Higher-order Forward Mode Automatic Differentiation as Weight Sharing
Here, we review higher-order forward mode, also known as Taylor- moydor-mode, automatic differentiation [19, 18, 3, tutorial in SCJ. Many PDEs only incorporate first- and second-orde 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 cully-connected and clement-wise activation layers. Our goal is to evaluate first-and second-order partial derivatives of the form 0,, u, 02 22,,,, u. for i, 1,..., d. At the
first layer, set (0) E Rdd, 0m., (0) = =e; E Rdd i.e., the i-th basis vector and 0 22,2, z 0 E Rd.
For linear layer fac (z(1-1))- = ( 2(1)(1110 (1-),, applying the chain rule yields the propagation rule
(1) = ( z (1-1) E 44
(5a)
0, 2 (1) = w 0z 2 (1-1)
E ER4
(5b)
02 a,
(1) g2 11-1)
E R
(5c)
2112,
The propagation rule through nonlinear element- activation layer 2 (1-1) + (2((111) is
2 (1) 1
(6a)
0z 0, (1) (((2(1-1) o Or x1-11
" (1-1) 1) e 0² 0" (1 @ 0r a (-11) + +z((((11)) S2.,, x (0-1)
(6b)
02 Gz,e, (1) 01
h0)
(6c)

---

Forward Laplacian For differential operators of special structure, we can fuse the Taylor-moo forward propagation of individual directional derivatives in Equations (5) and (6) and obtain 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 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
A 21) W Aaz
(7a)
Agz o' (1(2(11)) O Az 11-1) o" (z(21111)) @ 20-12021))62 E R
(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 more general second-ore order linear PDE operator L = =Eis=t 4.50 the forward pass for linear layer is (z(1) ( (z(1-1) E RA generalizing (7a), and similarly for Equation (7b)
Cz =o"=='(('(--)) (1- o (2"1)) + Σ
(2
On, 0n, (1-1) o Ony 20-1) E E
1-=1
SC.3 for details. This is different from [30], which transforms the input space such that the coefficients are diagonal with entries {0, 11,. reducing the computation to two forward Laplacians.
Importantly, the computation of higher-order derivatives for linear layers boils down to forward pass through the layer with weight sharing over the different partial derivatives (Equation (5)), and weight sharing can potentially be reduced depending the differential operator's structure (Equation (7a)). Therefore, we can use the concept 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 linear layer (suppressing Q in No)
e (0) - Σ (Jwa Un))) Jwan An
Because we made the Laplacian computation explicit through Taylor-mode autodiff ($3. specifically Equation (7a)), we can stack all output vectors that share the layer's weight into matrix z!!
xB with S = d + 2 and columns ( zn _ On, .!)) (!)
2,1+d Idd Or (1) and
z n,24d Aa (likewise - (1-1) E 1hal-1)S for the layer inputs), then apply the chain rule
s
Jw = (U
J)) Ag un
CRAR-1)
CRK()
which has structure similar to the Jacobian in $2.2, but with an additional sum the S shared vectors. With that, we can now express the exact interior Gramian for a layer as
S S
G n (
ΣV
(1-1) n,6"
((-1)
(1 (1) @ In..Inn''
(8)
N
Next, we want to approximate Equation (8) with a Kronecker product. To avoid introducing new convention, we rely on the KFAC approximation for linear layers with weight sharing developed by Eschenhagen et al. specifccil,,,,lyy, the approximation called KFAC-expand. This drops all terms with & s', then applies the expectation approximation from $2.2 over the batch and shared axes:
KFAC for the
matrix of a Laplace operator
N
N.S
0
~
7 (11)
(1
A 4
& B !
NS
N
5

---

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 R.M.
(u, u....... D
(10)
where tr collects all partial derivatives of order m. For E (o. there are Sm
(" (d-m--1) independent partial derivatives and the total number of independent derivatives is
S Em=g Sm -xx²). is smooth mapping from all partial derivatives to R
V: R +RM To construct PINN loss for Equation (10), feed the residual rg,n = RRM where D uo RdxSSm into smooth convex
V (uo (mn), Daugll (xn), criterion function P: RM - IR,
(0 = Σ ((ra.(()).
(11)
n=1
The generalized Gauss-Newton (GGN) matrix [51] is the Hessian of Ly (0) when the residual linearized w.r.t. 0 before differentiation. It is positive semi-definite and has the form
1 N
(
Σ (0))),(0)) A(raa) (Jorn..
(12)
with A(r) = V+((r))()))))) RMM the criterion's Hessian, e.g. e.g. {(r) 1/2/| r1 and A(r) IM-
Generalizing the second-order Taylor-mode from $3.1 to higher orders for the linear layer, we find
D = W (1) (t-1)
E ERA x
(13)
for any m. Hence, we can derive a forward propagation for the required derivatives where linear layer processes at most vectors", i.e. the linear layer's weight is shared over the ma-
trices x (1-11 = 2(1-1), (1=(1-1), (1-1) Stacking them into a matrix (1-1) =
(2(1-1), Dd= (1-1) ...., (1=1111)) 1 (and Z!! for the outputs), the chain rule yields
0 Σ
(Jwis (JwZz (J
A (rra,) (Jz.rr))) (Jwies 2,50)
N.S.S
Jwin 2!02))
(J2 rean)
A(rea) (J20), raam) (Jwis 2,00)
N
8,8,8
N.S.s
(1-1) ((-1) n,
J (Izل) ra
T
A A(ran)
Jz(e ren
N
n.
where -(--1)
Following the same steps in $3.2, apply the KFAC-expand approximation from [17] to obtain the generalization of Equation (9):
1)
E "-"
denotes the s-th column of Z (1-1)
KFAC for the GGN matrix of a general PDE operator
N.S
N.S Σ
~
NS
n.s
((-1))(11)17
1)T
N
J
200 ram
A(ram) (J22) raam) (14)
118=1
To bring this expression even closer to Equation (9), we can re-write the second Kronecker factor using an outer product decomposition A(ra,)) = -M h.mbom with Ln.m E RRM then introduce (1)
go.e.m : (J = r 4n (R) and write the second term as I/N E In.s.mgo.s.m. similar to the chonecker-factoree low-rank (KFLR) approach of Botev et al. [5].
(1) (i)
Depending the linear operator, may reduce weight sharing, demonstrated for the Laplacian in $3.1.

---

KFAC for variational problems Our proposed KFAC approximation is not limited to PINNs and can be used for variational problems of the form
min / E(u, Uuu, d
(15)
where L: R - IR is convex function. We can perceive this as a special case of the setting above with V 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 [16].
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. SB shows pseudo-code for our method on the Poisson equation.
At iteration 4, we approximate the per-layer interior and boundary Gramians using our derived Kronecker approximation (Equations (9) and (14)), G!l)
A
B and G!l) 002 = A 22 & B 0, f2 f1t
f2 ~
Exponential moving average and damping For preconditioning, we accumulate the Kronecker factors A B.! over time using an exponential moving average A!! 3A (1 A) of factor 3 E [0, (identically for B similar to the original KFAC. Moreover, we apply the same constant damping of strength À to all Kronecker factors, A. c A +XI and _ B +XI
such that the curvature approximation used for preconditioning at step t is
G..t ~ diag
(1) F2
(1)
ALL)
@ Bl.!) + diag 20, Dr
@
B 00 20
(L) 20,
@ Bly))
Gradient preconditioning Given layer ('s mini-batch gradient gi 0L(0:)/00 E Redd) we (1)
obtain update direction A - (A & B A (!)) ( ''1g E Reeh using the trick of [38, Appendix I] invert the Kronecker
an. )
At via dige-decoopooiin all Kronecker factors.
Learning rate and momentum From the preconditioned gradient A, E R,, we consider two different updates 0.+1 0, 8, we call KFAC and KFAC* KFAC uses momentum over previous updates, 8, = _ + A,, and μ chosen by the practitioner. Like ENGD, it uses a logarithmic grid line search, selecting 6, = a. with a, = arg min & L(0, + (ε;;) where a E [2-30....,,) 20)). KFAC* uses the automatic learning rate and momentum heuristic of the original KFAC optimizer. parametrizes the iteration's 's update as 10,1(a, 11) a , + μóz then obtains the optimal parameters by minimizing the quadratic model m(6+11) = L(O,) + 5T + 1/2ōTT. (G(O.) + + with the exact damped Gramian. The optimal learning rate and momentum arg min,,, m( are
(ar) -- -- A, AT + +
A, + AA, & A G(O.) + AAY 5, 8, G(O.) + |
-1
A 6, 9:
(see [38, Section 7] for details). The computational cost is dominated by the two Gramian-vector products with A, and 8,. By using the Gramian's 's outer product structure [12, 45], we perform them with autodiff [48, 51] using one Jacobian-vector product each, recommended in [38].
Computational complexity Inverting layer l's Kronecker approximation of the Gramian requires Ohh)) + 1 time and !(h n1(+1)2) storage, where (1) the number of neurons in the (-th layer, whereas inverting the exact block for layer would require (110) (((1)) time and (), (((+)) ((+1)) memory. In general, the improvement from the Kronecker factorization depends on how close square the weight matrices of layer are, and therefore the architecture. In practise, the Kronecker factorization usually significantly reduces memory and run time. Further improvements can be achieved by using structured Kronecker factors, e.g. (block-Jdiagonal matrices [32].
We use the forward Laplacian framework in our implementation, which we found to be significantly faster and more memory efficient than computing batched Hessian traces, see $C.4.
7

---

D = 257
D = 9,873
D = 116,097
10
L error
10
10
10-3
10-3
10
10
10"5
10-5
10"
10
10
10a
10a
10
10a
10a
10"
10'
10a
10a
Time
[s]
Time [s]
Time [s]
SGD
Adam
Hessian-frec
LBFGS
ENGD (full)
ENGD (layer-wise)
KFAC
Figure/ 1: Performance of different optimizers on the 2d Poisson equation (16) measured in relative L2 error against wall clock time for architectures with different parameter dimensions D.
D 449
D 10,065
D 116,865
10
10
10
L2 error
10-3
10-2
10
10
10
10-4
10
10
10a
10a
10a
10
10a
10a
10
10
10a
10a
Time [s]
Time [s]
Time [s]
SGD
Adam
Hessian-froc
LBFGS
ENGD (full)
ENGD (laycr-wise)
KFAC
KFAC
Figure 2: Performance of different optimizers on the (4+1)d heat equation (17) measured in relative
_ error against wall clock time for architectures with different parameter dimensions D.
4 Experiments
implement KFAC, KFAC*, and ENGD with either the per-layer or full Gramian in PyTorch [46]. As 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 fully-featured implementation from Tatzel et al. [55] which offers many additional hcuristics like adaptive damping, CG backtracking, and backtracking line search, allowing this algorithm to work well with little Hyper-paramette tuning, As basclines, 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 dyper-parameter protocol). For aandom/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 L2 error estimated on held-out data set with the known solution a compute cluster with RTX 6000 GPUs (24 GiB RAM) in double precision, and we use the same computation time budget for all optimizers hyper-paran parameters,, as well as training curves over iteration count rather than time, are in $A.
using Weights & Biases [60] (see $A.1 1 for the exact
the studied PDE. To be comparable, all runs are executed on
fixed PINN problem. All search spaces and best
Pedagogical example: 2d Poisson equation We start with low-dimensional Poisson equation from Müller & Zeinhofer [41] to reproduce ENGD's performance (Figure 1). It is given by
Au(r, y) = 2n sin (T) sin(Try) for (x. E [0
(16)
u(x, v) = for (I, v) E 000
We choose 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, shallow and two with five fully-comnected layers 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-ENGr for long

---

5d Poisson, D 116,865
10d Poisson, D 118,145
100d Poisson, D 1,325,057
10a
10
L error
10
10
10-3
10
10
10"
10
10a
10a
10"
10a
10a
10a
10a
10a
10
Time [s]
Time [s]
Time [s]
SGD
Adam
Hessian-free
LBPGS
Figure 3: Optimizer performance on Poisson equations in high dimensions and different boundary conditions measured in relative L2 error against wall clock time for networks with D parameters.
phase, outperforming the first-order and quasi-Newton baselines. In contrast to ENGD, which runs out of memory for networks with more than 10000 parameters, KFAC scales to larger networks (Figure I, 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 fig. A10.
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 four-dimensiona heat equation
O O, x) x K ut, x) = 0 for t E [0, 1], x e [0,
u0, x) sin (2x;)
for x E [0, 1J",
(17)
u(t, x) exp(-t) Σ sin() for E [0, 1]. E 000,
with diffusivity constant 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 $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 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 additional experiment for the MLP with D 116 864 where we tune batch sizes, batch sampling frequencies, and all hyper-paree parameters with generous search spaces using Bayesian search ($A.10) 10). We find that this does not significantly boost performance of the other methods (compare Figures 2 and A14). 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 the same architectures, see $A.3 and fig. A7.
High-dimensional Poisson equations To demonstrate scaling to nigh-dimensional PDEs and even larger neural networks, we consider three Poisson equations (d = 5, 10, 100) with different boundary conditions used [16, 41], which admit the solutions
W. (w)
coS()T for x E [0, [0,1]
(18)
u, (x:)
F2k- for a E [0, 1]10
n, (x:) = ||0:|| for x E [0, 100
100
We use the same architectures as before, but with larger intermediate widths and parameters up to million (Figure 3). Due to lacking references for training such high nigh-dimensional problems,

---

select all hyper- hhper-paramenn parameters via Bayesian search, including batch sizes and batch sampling frequencies (details in $A.5). We see 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 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 (details in $A.4). In this experiment, KFAC also achieved similar performance than Hessian- and outperformed SGD, Adam, and LBFGS (Figure A8).
we conduct a
(9+1)d Tokker-Plance equation To show the applicability to nonlinear PDEs, we consider a Fokker-Planne equation in logarithmic space. PINN formulations of the Pokker-Plance equation have been considered in [23, 54]. Concretely, we are solving nine-dimensionn equation of the form
d Va(t,)) x) a (V x)1 Aq x) = O, q()) = log(p (0)). (19)
deg(t, x) - 2 2
with d = 9, [o, 1] and x R.. where in practice R is replaced by [-5, 5]9. The solution gs = log(p") and p" is given by p" (t, We model the solution with medium sized tanh-activateds MLP with = 118 145 parameters, batch sizes N₂ 000, Na 1 000, and we assign each run a computation time budget of 6 000 s. As in previous experiments, the batches sampled every iteration for all optimizers except for KFAC and KFAC*, which use the same batch for ten steps (details in $A.1 1). Figure 4 reports the 2 error over training time. Again, KFAC is among the best performing optimizers offering competitive performance to Hessian-free and clearly outperforming all first-order methods.
10
(e((-exp(-t))21
L crror
SGD
100
Adam
Hessiam-free
LBFGS
10-1
KFAC
KFAC*
re-
mm)
10a
10
10a
Time [s]
Figure 4: Performance of different optimizers on (9+1)d logarithmic Fokker Planckk equation in relative L2 eгror against wall clock time.
5 Discussion and Conclusion
We extended the concept of Kronecker-factoree approximate curvature (KFAC) to Gauss-Newrr matrices of Physics-infee informed neural network (PİNN) 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 's compute graph as a forward net with shared weights, then apply the eecently-develped formulation of KFAC for linear layers with weight sharing. Empirically, we find that our KFAC-based optimizers are competitive with expensive second-orde methods on small problems and scale to high-dimensions 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 apply to arbitrary architectures and PDES, as described in $3.3. We are excited that our current algorithms show promising performance when compared second-order methods with sophisticated heuristics. In fact, the original KFAC optimizer itself [38] relies heavily on such heuristics that are said 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 's performance can further improved, e.g. by exploring
derive KFAC (Taylor-mode, weight sharing)
(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 [31], (4) using structured Kronecker factors [32], (5) computing the Kronecker factors 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].
10

---

Acknowledgments and Disclosure of Funding
The authors thank Runa Eschenhagen for insightful discussions on KFAC for linear weight sharing layers. FD would like to thank Luca Thiede for his adamant questions about Taylor mode and forward Laplacians. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. JM acknowledges funding by the Deutsche Forschungsgemeinschaff (DFG, German Research Foundation) under the project number 442047500 through the Collaborative Research Center Sparsity and Singular Structures (SFB 1481). MZ acknowledges support from an ETH Postdoctoral Fellowship for the project "Reliable, Efficient, and Scalable Methods for Scientific Machine Learning".
References
[1] Amari, S.-I. Natural gradient works efficiently in learning. Neural computation, 10(2)251-277. 1998.
[2] Benzing, F. Gradient Descent on Neurons and its Link to Approximate Second-order Optimiza- tion. In International Conference Machine Learning (ICML), 2022.
[3] 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 (NeurlPS); Workhop Program Transformations for ML, 2019.
[4] Bonfanti, A., Bruno, G., and Cipriani, C. The Challenges of the Nonlinear Regime for Physics- Informed Neural Networks. arXiv preprint arXiv:2402.03866 2024.
[5] Botev, A., Ritter, H., and Barber, D. Practical Gauss-Newtor optimisation for deep learning. In International Conference on Machine Learning (ICML), 2017.
[6] Bottou, L.,, Curtis, F. E., and Nocedal, J. Optimization methods for large-scale machine learning. SIAM Review (SIREV), 60, 2016.
[7] 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.
[8] Clarke, R. M., Su, B., and Hemández--obbtto, J. M. Adam through second-order lens. 2023.
[9] 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.
[10] Dangel, F., Harmeling, S., and Hennig, P. Modular Block-diagonal Curvature Approximations for Feedforward Architectures. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
[11] Dangel, F,, Kunstner, F., and Hennig, P. BackPACK: Packing more into Backprop. In Interna- lional Conference on Learning Representations (ICLR), 2020.
[12] Dangel, F., Tatzel, L., and Hennig, P. VIVIT: Curvature access through the generalized gauss- newton's ' low-rank structure. Transactions on Machine Learning Research ('MMLR), 2022.
[13] Daw, A., Bu, J., Wang, S., Perdikaris, P., and Karpatne, A. Rethinking the importance of sampling in thysics-informen neural networks. arXiv preprint arXiv: 2rXiv:2207.023388 2022.
[14] 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.
[15] Dissanayake, M. and Phan-Thien, N. Neural hetrllewttwoo-bbasee approximations for solving partial differential equations. communications in Numerical Methods in Engineering, 10(3):195-201 1994.
[16] 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.
11

---

[17] Eschenhagen, R., Immer, A.. Turner, R. E., Schneider, F., and Hennig, P. Kronecker-Facthe Approximate Curvature for Modern Neural Network Architectures. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[18] Griewank, A. and Walther, A. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
[19] Griewank, A., Juedes, D., and Utke, 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-16) 1996.
[20] Grosse, R. and Martens, J. Kronecker-Factoree Approximate Fisher Matrix for Convolution Layers. In International Conference on Machine Learning (ICML), 2016.
[21] 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.
[22] Heskes, T. On "natural" learning and pruning in multilayered perceptrons. Neural Computation, 12(4):881-901, 2000.
[23] Hu, Z., Zhang. Z., Karniadakis, G. E., and Kawaguchi, K. Score-based physics-informed neural networks for high-dimensional fokker-planck equations. arXiv preprint arXiv:2402.07465, 2024.
[24] Jnini, A., Vella, F., and Zeinhofer, M. Gauss-Newton Natural Gradient Descent for Physics- Informed Computational Fluid Dynamics. arXiv preprint arXiv:2402. 10680, 2024.
[25] Johnson, M. J., Bettencourt, J., Maclaurin, D., and Duvenaud, D. Taylor-made higher-order au- tomatic differentiation. 2021. URL https://pg thub. com /ttg/////aa///ttl///////// jet.pdf. Accessed January 03, 2024.
[26] 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, 14:26554--26565 2021.
[27] 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.
[28] Lagaris, E., Likas, A., and Fotiadis, D. Artificial neural networks for solving ordinary and partial differential equations. IEFF transactions on neural networks, 955)9887-1000, 1998.
[29] 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.
[30] 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 Al4DiiferentialEquation In Science, 2024.
[31] Lin, W., Duruisseaux, V., Leok, M., Nielsen, F., Khan, M. E., and Schmidt, M. Simplifying Momentum- based Riemannian Submanifold Optimization. 2023.
[32] Lin, W., Dangel, F., Eschenhagen, R., Neklyudov, K., Kristiadi, A., Turner, R. E., and Makhzani,
A. Structured inverse-free natural gradient descent: Memory-efficient & numerically-stable KFAC. In International Conference on Machine Learning (ICML), 2024.
[33] Liu, S., Su, C., Yao, J., Hao, Z. Su, H., Wu, Y., and Zhu, J. Preconditioning for physics-infommed neural networks. arXiv preprint arXiv:2402.0053/ 2024.
[34] Lu, L., Meng, X., Mao, Z., and Karniadakis, G. E. DeepXDE: A deep learning library for solving differential equations. SIAM Review, 63(1):208-22) 2021.
[35] Markidis, The old and the new: Can physics-informed deep-learning replace traditional linear solvers? Frontiers in big Data, 4:669097, 2021.
12

---

[36] Martens, J. Deep learning via Hessian-free frre optimization. In International Conference on Machine Learning (ICML), 2010.
[37] Martens, New insights and perspectives on the natural gradient method, 2020.
[38] Martens, and Grosse, Optimizing Neural Networks with Kronecker-factoree Approximate Curvature. In International Conference on Machine Learning (ICML). 2015.
[39] Martens, J., Ba, J.. and Johnson, M. cronecker-factorie Curvature Approximations for Recur- rent Neural Networks. In International Conference on Learning Representations. 2018. URL https://openrewe
net hetefoormm/ee www.oooioiiiio.ooo... Hoy.iiii-yyHtttAAAba
[40] 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.
[41] Müller, and Zeinhofer, M. Achieving high accuracy with PINNs via energy natural gradient descent. In International Conference Machine Learning, pp. 114411254855... PMLR, 2023.
[42] Müller, J. and Zeinhofer, M. Optimization in SciML-A Function Space Perspective, arXiv preprint arXiv:2402.07316 2024.
[43] Nabian, M. A., Gladstone, J., and Meidani, Efficient training of physics-informed neural networks via importance sampling. Computer-Aided Civil and Infrastructure Engineering, 36 (8):962 977, 2021.
[44] 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.
[45] Papyan, Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet Hessians. In International Conference on Machine Learning (ICML), 2019.
[46] 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., De Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L.., Bai, J., and Chintala, S. PyTorch: An Imperative Style, Perforaere Deep Learning Library. Advances Neural Information Processing Systems (NeurIPS). 2019.
[47] 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. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14, 2021.
[48] Pearlmutter, B. A. Fast Exact Multiplication by the Hessian. Neural Computation, 1994.
[49] Petersen, F., Sutter, T., Borgelt, C., Huh, D., Kuehne, H.,, Sun, Y., and Deussen, ISAAC Newton: Input-based Approximate Curvature for Newton's 's Method. In International Conference on Learning Representations (ICLR), 2023.
[50] Raissi, M., Perdikaris, P., and Karniadakis, G. Physics-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
[51] Schraudolph, N. N. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation, 2002.
[52] Sirignano, J. and Spiliopoulos. K. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1364 2018
[53] Skorski, Chain nules for hessian and higher derivatives made easy by tensor calculus. arXiv preprint
2019.
[54] Sun, J., Berner, J., Richter, L., Zeinhofer, M., Müller, J., Azizzadenesheli, K., and Anandku- mar, A. Dynamical measure transport and neural pde solvers for sampling. arXiv preprint arXiv:2407.07873, 2024.
13

---

[55] Tatzel, L., Hennig, P., and Schneider. F, Late-Phase Second- Order Training, In Advances in Neural Information Processing Systems (NeurIPS), Workshop Has Trained Yet?, 2022.
[56] 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.
[57] Wang, S., Teng, Y., and Perdikaris, P. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal Scientific Computing, 43(5):A3055-A3333 A3081, 2021.
[58] Wang, S., Sankaran, S., and Perdikaris, P. Respecting causality is all you need for training physics-informed neural networks. arXiv preprint arXiv:2203. 074404 2022.
[59] Wang, S., Yu, X., and Perdikaris, P. When and why PINNs fail to train: neural tangent kernel perspective. Journal of Computational Physics, 449:1 10768, 2022.
[60] Weights and Biases. Experiment Tracking with Weights and Biases, 2020. URL https: / //www. wandb. ai/. Software available from wandb.ai.
[61] 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: 15671, 2023.
[62] Zampini, S., Zerbinati, U., Turkiyyah, G., and Keyes, D. PETScML: Second-order solvers for training regression problems in Scientific Machine Learning. arXiv preprint 1281: 2024.
[63] 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.
[64] Zeng, Q.. Kothari, Y., Bryngelson. S. H., and Schäfer, F. Competitive physics informed networks. arXiv preprint arXiv:2204. //14: 2022.
14

---

Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks (Supplementary Material)
Experimental Details and Additional Results
15
A.1 Hyper-Parameter Tuning Protocol
15
A.2 2d Poisson Equation
16
A.3 5d Poisson Equation
20
A.4 10d Poisson Equation
22
A.5 5/10/100-d Poisson Equations with Bayesian Search
................................
24
A.6 PINN Loss for the Heat Equation
................................
27
A.7 1+1d Heat Equation
28
A.8 4+1d Heat Equation
31
A.9 Robustness Under Model Initialization for 4+1d Heat Equation
34
A.10 4+1d Heat Equation with Bayesian Search
A.11 9+1-d Logarithmic Pokker-Planck Equation with Random Search
36
Pseudo-Code: KFAC for the Poisson Equation
38
C Taylor-Mode Automatic Differentiation & Forward Laplacian
39
C.1 Taylor-Mode Automatic Differentiation
39
C.2 Forward Laplacian
40
C33 Generalization of the Forward Laplacian to Weighted Sums of Second Derivatives..
42
C.4 Comparison of Forward Laplacian and Autodiff Laplacian
43
Backpropagation Perspective of the Laplacian
43
D.1 Hessian Backpropagation and Backward Laplacian
44
D.2 Parameter Jacobian of the Backward Laplacian
46
D.3 Gramian of the Backward Laplacian
46
A Experimental Details and Additional Results
A.1 Hyper-Parametter Tuning Protocol
In all our experiments, we tune the following optimizer hyper-parameters and otherwise use the PyTorch default values:
SGD: learning rate, momentum
15

---

Adam: learning rate
type of curvature matrix (Hessian or GGN), damping, whether to adapt damping over time (yes no), maximum number of CG iterations
Hessian-free:
LBFGS: learning rate, history size
ENGD: damping, factor of the exponential moving average applied to the Gramian, initial- ization 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- paremeeaaes.rs. Each individual run is executed in double precision and allowed to run for given time budget, and we rank runs by the final L2 error on 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 based approach. First, we choose a relatively wide search space and limit to approximately 50 runs. space based the first round, then re-run for another approximately 50 runs. We will release the details of all implementation.
second round, we narrow down the syper-paramente
search spaces, as well as the hyper-parameteet for the best runs in our
A.2 2d Poisson Equation
Setup We consider a two-dimensional Poisson equation - Au((x, 2x sin(T.) sin(Ty) on the unit square (x, y) [0, 1|2 with sine product right-hand side and zero boundary conditions u(x, y) = 0 for (x, y) 0 1ja. We choose a single set training points with Nºp 900, Nan = 120. The L2 error is evaluated on separate set of 000 data points using the known solution w. (,, 3) = sin (TZ) (T Each run is limited to a compute time of 000 We compare three MLP architectures of increasing size, each of whose linear layers Tanh-activated except for the final one: a shallow 2 + 64 + MLP with D = 257 trainable parameters, a five layer 48 - 48 - 1 MLP with D 873 trainable parameters, and five layer 2 - 256 > 256 + 128 - 128 + 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 tested in the experiments. Figure A5 visualizes the results, and Figure A6 illustrates the learned solutions over training for all optimizers the shallow MLP.
2 - 64 - 64
Best run details
The runs shown in Figure A5 correspond to the following hyper-parameters:
2 + 64 MLP with D 257
SGD:
learning rate: 1.805 015 10-2, momentum: 9.9 10-1
Adam: learning rate: 1.692 339- 10-3
Hessian-free: curvature matrix: GGN, initial damping: 500, constant damping: no,
maximum CG iterations: 300
LBFGS: learning rate: 5 10-1, history size: 150
ENGD (full): damping: 10-10, exponential moving average: 3 10-1, initialize
Gramian to identity: yes
ENGD (layer-wise): damping: exponential moving average: 9 10-1, initialize
Gramian to identity: yes
KFAC: damping: 1.544099-11
10 momentum: 5.117 -0-1 exponential moving
average: 4.496 490 10-1 initialize Kronecker factors to identity: yes
KFAC*; damping: 1.215 640 10-10 exponential moving average: 9.263 314 10-1,
initialize Kronecker factors to identity: yes
2 + 64 64 + 48 - 48 - MLP with D = 9.8
-GGD: learning rate: 3.758 303 10-3, momentum: 9 10
Adam: learning rate: 2.052 448 1
16

---

(a)
D 257
9,873
116,097
10
10-1
L error
10-2
10-3
10-9
10a
104
109
2
Loss
10-3
10
10-7
10-10
109
10
10a
10a
100
10
10
10a
10
10
10a
(b)
10a
10
10-1
Lg error
10-2
10-3
10-3
10
1073
10-5
10a
109
10a
10-2
Loss
10-3
10
10-7
10-10
10
10a
10
10'
10a
10a
10a
Iteration
Iteration
Iteration
SGD
Adam
Hessian-free
LBFOS
ENGD (full)
ENGD (layer-wise)
KFAC
KFAC
Figure A5: Training loss and evaluation L2 error for learning the solution to 2d Poisson equation over (a) time and (b) steps. Columns are different neural networks.
Hessian-free: curvature matrix: GGN, initial damping: 10-1, constant damping:
maximum CG iterations: 350
LBFGS: learning rate: 1 10-1, history size: 200
ENGD (full): damping: 1 10- 10, exponential moving average: 10-1, initialize
Gramian to identity: no
ENGD (layer-wise): damping: 1-10-6. exponential moving average: 3- 3-10-1, initialize
Gramian to identity: no
KFAC: damping: 2.640 1000-100111 momentum: 9.995 595- 105 exponential moving
average: 5.556 664 10 initialize Kronecker factors to identity: yes
KFAC*; damping: 2.989 247 10-13, exponential moving average: 6.258 340 10-1,
initialize Kronecker factors to identity: yes
2 + 256 + 256 + 128 + 128 - MLP with D = 116097
SGD: learning rate: 2.478 674 10-3, momentum: 9 10
Adam: learning rate: 6.406 108 10-4
Hessian-frre:: curvature matrix: GGN, initial damping: 50, constant damping: no,
maximum CG iterations: 350
LBFGS: learning rate: 1 1 history size: 225
17

---

Optimizer
First step
0.1 % trained
1% trained
10% trained
True solution
SGD
Adam
LBFGS
Hessian-free
ENGD (full)
ENGD (layer-wise)
KFAC
KFAC*
Figure A6: Visual comparison learned and true solutions while training with different optimizers for the 2d Poisson equation using two-layer MLP (corresponding the curves in Figure left). All functions are shown on the unit square (x, y) E = [0; 1]2 and normalized to the unit interval.
KFAC: damping: 1.710 16910229-111 200-109-110 13 momentum: 8.484 996-999-1 10-1, exponential moving
average: 9.636 460 10-1. initialize Kronecker factors identity: yes
KFAC*; damping: 1.232 407 10-13, exponential moving average: 9.8 207 10-1, initialize Kronecker factors identity: yes
Search space details The runs shown in Figure A5 were determined to be the best via search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search (l4 denotes uniform, and Ll log-uniform distribution):
. 2 2 64 MLP with D 257
18

---

SGD: learning rate: [ 10 1013; 3,1 1)" momentum: 4((0, 1..00"06.00'..... 10.1,9 10-1 9.9 10 1))
Adam: learning rate: Lu((5 10 10-4; 10-0011)
Hessian-free; curvature matrix: U((GGN, Hessian)), initial damping: U 1,
10-2.1 1. 10000,1111 10))))))), constant damping: U( (no, yes)), maximum CG iterations:
((550, 250})
LBFGS: learning rate: [1 10-1,2211000 10-1, 11.111..111......1........11111.
10-1, ..
1002,11-22200 2- 10-2, 1. 1- 2,1-10-3) 10-
history size: U((75, 100, 125, 150, 200, 225,
ENGD (full): damping: W((1
1.1000-11 10012, (00)).0)), exponential 10-110011,-110,0000 10- initialize Gramian to
1..10-,.11100
moving average: (((0,33 3 -0
identity: ((no, yes})
ENGD (laeerwwise): damping: [((11100( 1100 1-10-10 0}). exponential moving average: (0, 10 10-66---1-1 10001,9-10-1,00 10 10)1)) initialize Gramian to
-
identity: 2/ [no, yes})
KFAC: damping: [1((1110 13, 1011--117)) "),, momentum: 21[0; 9.9- 10-1101-)) (). exponential
moving average: 24( 9.9 10 '1) initialize Kronecker factors to identity: yes
KFAC*; damping: [[[( [[( 10-13;1 1. 10 1007)), exponential moving average: U([0;
10 10-11 'I), initialize Kronecker factors to identity: yes
2 + 64 64 48 + 48 - MLP with D 9
SGD: learning rate: [u( 10 10-3; 1. -1), momentum: 3((0,3- 10- 1611,,9 10-1))
Adam: learning rate: [u((1 10 1::4::::, 10=
Hessian-free; curvature matrix: U(IGGN, Hessian)), initial damping: U({1,
10-',1-101 2,1 1,1-1003, 10-3,1-11111 1,1-1004) constant damping: U((no, yes)), maximum CG
iterations: ((50 250))
LBFGS: learning rate: L((5 1005-10111,0 1.-
1111101--00 5 10 10-2,1222220 2,2 10-2, 2, 10- -1)-)),
history size: ((550, 75, 100, 125, 150, 200,
- ENGD (full): damping:
11 1 10-12 10-12,03 ,0)), 10) initialize Gramian exponential moving average: 10 10-
10
10
1
10-11,
10
to identity: [nnoo, yes))
ENGD (layer-wise): damping: U((1-10-2, 10
1- 10 14- 10 1-11--1011 10-")), 10-1110_,611--0000 10111101100000 1,01,99
exponential moving average:
9.9 10 10-1)), initialize
Gramian to identity: U(noo yes))
KFAC: damping: [U((1 10 10-13;1 10-7)), momentum: u( 9.9 10 10-1I), expo-
nential moving average: U([0; 9.9 10
10-1)) IJ), initialize Kronecker factors identity:
[(noo, yes))
KFAC*; damping: [u([1- 10-13; 10 7),, exponential moving average: U(0;99.
10- initialize Kronecker factors to identity: [1([no,
2 256 > 256 + 128 + 128 MLP with D 116 097
SGD: learning rate: [u 10 3,1 1- 10 -0), momentum: 10-1))
1.-1,,9
Adam: learning rate: [u((1 10 10-4;5 4::5-10'1)
Hessian-free: curvature matrix: U u({GGN, Hessian)), initial damping: U({1, 10-1.1 10-2,1 2,1 10 10-4)), constant damping: y(sno, maximum CG
iterations: (/50, 250))
LBFGS: learning rate: 2/((5 10-1,2.11 10-1.11 10-1,5 10- 10 10-2,11 10 (0-2)),
history size: (1500, 75, 100, 125, 150, 175, 10,,,00,,22))
KFAC: damping: LM(( 13 -10-7)). momentum: 9.9 10-1)), expo-
nential moving average: U((00 9.9 10-1]), initialize Kronecker factors identity: u([no,yes]
KFAC*; damping: LU((1 13; 1- exponential moving average: U([ 9.9.
10-7)),
10-41), initialize Kronecker factors to identity: yes))
19

---

A.3 5d Poisson Equation
Setup We consider a five-dimensional Poisson equation -A((x) n² Ei= cOs( on the five-dimensional unit square [0, 1] with cosine sum right-hand side and boundary conditions u(x) = Σi=1 (T2( for æ E 0 1js. We sample training batches of size No = 3 000, Nan = 500 and evaluate the L2 error on separate set of 30 000 data points using the known solution (x) Ei=1 cos(xT;). All optimizers except for KFAC sample
training batch each iteration. KFAC only re- re-samples every 100 iterations because we noticed significant improvement with multiple iterations fixed batch. To make sure that this does not lead an unfair advantage of KFAC, 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 consistent with this additional experiment (compare the rightmost column of Figure A7 and the leftmost column of Figure A9). Each run limited s.00 We compare three MLP architectures of increasing size, each of whose linear layers are Tanh-activated except for the final one: a shallow 5 + 64 + 1 MLP with / = 449 trainable parameters, five layer 5 + 64 64 + 48 + 48 MLP with D 10 065 trainable parameters, and five layer - 256 + 256 - 128 + 128 + MLP with D = 116 864 trainable
parameters. For the biggest architecture, full and layer-wise ENGD lead -ettoofymmrry errors and are thus not tested in the experiments. Figure A7 visualizes the results.
(a)
D) 449
10,065
116,865
10a
10
Lg error
10"2
10-1
10
10
10-3
10
THM
10a
101
Loss
10-3
10
109
10
10-5
10-4
mm
10a
10a
10a
b
10a
10a
Li error
10:
10-1
10
3
10-3
10a
10a
Loss
10-3
10
10"
10
10
10-4
10
30
10a
1
10
10
10a
Iteration
Iteration
Iteration
SGD
Adam
Hessiam free
LBFGS
ENGD (full)
ENGD (layer wise)
KFAC - KFAC*
Figure A7: Training loss and evaluation Le error for learning the solution to 5d Poisson equation over (a) time and (b) steps. Columns are different neural networks.
20

---

Best run details
The runs shown in Figure A7 correspond to the following seper-parameters.
5 64 MLP with 449
SGD:
learning rate: 4.829 757 10-3, momentum: 9 10
Adam: learning rate: 1.527 101 10-3
Hessian-free: curvature matrix: GGN, initial damping: 5, constant damping: no,
maximum CG iterations: 350
LBFGS: learning rate: 5 10-2, history size: 125
ENGD (full): damping: 1 10-11 exponential moving average: 6 10-1, initialize
Gramian to identity: yes
ENGD (layer-wise): damping: 1-10-b, exponential moving average: 6- 6-10 initialize
Gramian to identity: yes
KFAC: damping: 3.030 734-10-11 10 momentum: 4.410 155- 1001 exponential moving
average: 3.260 663 10
initialize Kronecker factors to identity: yes
: damping: 9.835 853 10-13 exponential moving average: 28714287 10-1,
initialize Kronecker factors identity: yes
5 64 64 48 48 ML. with D 10 065
SGD: learning rate: 1.007 555 10-3 momentum: 9 10-1
Adam: learning rate: 6.999 994 10
Hessian-free: curvature matrix: GGN, initial damping: 1 10 constant damping:
no, maximum CG iterations: 350
LBFGS: learning rate: 1 10-1, history size: 225
ENGD (full): damping: 1 10
exponential moving average: 3 10-1, initialize
Gramian to identity: no
ENGD layer-wise): damping: 1-10-6, exponential moving average: 3 3-10-1, initialize
Gramian to identity: no
KFAC: damping: 1.183 063-10-12
momentum: 9.058 900 1000 exponential moving
average: 9.588 846 10-1 initialize Kronecker factors to identity: yes
KFAC*; damping: 2.829 461
10-10, exponential moving average: 9.001 393 10-1,
initialize Kronecker factors identity: yes
5 + 256 + 256 + 128 128 MLP with D = 116 865
SGD: learning rate: 1.924 173 10-3, momentum: 9 10
Adam: learning rate: 5.416 376 10-4
Hessian-free: curvature matrix: GGN, initial damping: 5, constant damping: no,
maximum CG iterations: 350
LBEGS: learning rate: 2 -00-2, history size: 225
KFAC: damping: 1.844 213- (0- momentum: 7.528 559-10" 10-1, exponential moving
average: 9.307 849 10-1 initialize Kronecker factors identity: yes
KFAC*; damping: 2.183 605 10 10 exponential moving average: 9.563 992 10-1,
initialize Kronecker factors to identity: yes
Search space details The runs shown in Figure A7 were determined to be the best via search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search (24 denotes a uniform, and LU a log-uniform distribution):
* 5 - 64 + MLP with D = 449
SGD: learning rate: [U[ 10 3, -1), momentum: 3 10-1,9 10-1))
Adam: learning rate: LM(11 1: -00-111
Hessian-frre: curvature matrix: u({GGN, Hessian)), initial damping: U((1, 10-1.1 10-2,1 2.1 10-3, 1- constant damping: (({n),
yes?), maximum CG
iterations: ((50 250))
LBFGGS: learning rate: (1(5 10-1,2 10-1, 1 10-1,5-10-2,, 10-2, 2, -,1110-21
-102)),
history size: ud 75, 100, 125, 150, 175, 200, 2253)
21

---

ENGD (full): damping: u([1 ([[[1-10-.. 1....110.1000000 10
10
10
1 10 10-11, 1 10012,000
exponential moving average: U([0,3 10-.3--00-,0000 6 to identity: ([(no,
10
9
10
1), initialize Gramian
ENGD (layer-wise): damping: 4[[1- 10
1 10- 11 10 10-5, ..1.111111... 10.1--10))) 10-
exponential moving average: Ud (0, 10- 10 10 ,.9. 10 10-1)), initialize
9.9
Gramian to identity: [[nno, yes))
KFAC: damping: [ 11, 11111001] 1)) momentum: ([; 9.9-10- exponential
moving average: (1 9.9 10 10- initialize Kronecker factors to identity: yes
KFAC*; damping: [U((1 10:18; 1..1111 10 7), exponential moving average: ((;
ID), 10-"1),
10
initialize Kronecker factors to identity: yes
5 + 64 + 64 + 48 - 48 + MLP with D = 10065
SGD: learning rate: cu((1 10 ;1 1 10 3)), momentum: 10-1))
Adam: learning rate: [U((1 10 4; 1)=
Hessian-free; curvature matrix: U({GGN, Hessian)). initial damping: U((1, 10-1,1 10-2, 10-3 10003,1-1--11 }), constant damping: (([no, yes)), maximum CG
iterations: (4(50, 250})
1001. LBFGS: learning rate: 2/(55 10-100010000 10001,2.00--111 10-1
5 10 2, 10-100- 1.1-10-31 10-2)),
history size: (150,, 75, 100, 125, 150, 175, 200, 220))
10 1 10011 1 10-12,00 0}), exponential moving average: U((0,3 10-6-0011.00 to identity: U( yes})
ENGD (full): damping: U({1 10-08,111-0
10-1)) initialize Gramian
ENGD (layer-wise): damping: u([1-
10-4, 10-10- 1--1--1011 10-")), 9.9 10- 10-11). initialize
exponential moving average:
60-10-6-111000 190011001911000 10- 1.9.
Gramian to identity: L4 (( yes))
KFAC: damping: [2((11 10110-0-11-111).. momentum: 4{ 9.9 9.9-10 10-1-))) exponential
1-100)]).
moving average: 44 10-1)), initialize Kronecker factors identity: yes
KFAC*; damping:
10-12; 10-12;1.10-91 10-
exponential moving average: [[[00) 9.9
1), initialize Kronecker factors to identity: yes
5 + 256 > 256 + 128 + 128 MLP with D = 116 865
SGD: learning rate: [ 10-3; 1 10 10-2)), 10-1))
21), momentum: U((0,3- 10-1, 10-1,0
Adam: learning rate: [u([1 10 10-4;5 10--))
Hessian-free; curvature matrix: U({GGN, Hessian)),, initial damping: U({1, 10-1,1 10-2,1 122111111,1111110000011100011111 10-3,1
100-31111111)11 10-4)), constant damping:
maximum CG
iterations: (((50, 250})
LBFGS: learning rate: 24(55 20-1--110-1 1 10-1011,,0 1,- - 10 2,2 10 100-1,1-110 2, 100 ()))))
history size: 4({50, 75, 100, 125, 150,
KFAC: damping: [24(11-0 11:1000011111111 momentum: 24[[0; 9.9 10 'I), exponential
moving average: L<[[; 9.9 10-' initialize Kronecker factors to identity: yes
KFAC*; damping: [ 10-12;1 -10-12;1110-)),, exponential moving average: L[[00 9.9
II), initialize Kronecker factors to identity: yes
A.4 10d Poisson Equation
Setup We consider 10-dimensional Poisson equation -Au(x) 0 on the 10-dimensional unit square æ E [0, 1 with zero right-hand side and harmonic mixed second order polynomial boundary conditions u(x) = Σ" 1 for 0[0,1]4. We sample training batches of size Ng = 3 000, Non = 1000 and evaluate the L2 error on a separate set of 30 000 data points using the known solution u. (z) = >%/ I2i-1 All optimizers except for KFAC sample new training batch each iteration. KFAC only re-samples every 100 iterations because we noticed significant improvement with multiple iterations
i=1
a fixed batch. Each run is limited to s000. We use 10 > 256 + 256 -> 128 - 128 + MIP with D 118 145 MLP whose linear layers are Tanh-activated except for the final one. Figure A8 visualizes the results.
22

---

(a)
10"
10a
10-1
10-2
L error
10-2
SGD
10-4
Adam
10
Hossian-free
10-8
LBFGS
10-4
10
100
10a
10a
10a
100
10a
10a
Time [s]
Time
(b)
100
10a
10-1
10-2-
Lg error
MMMMA
10-2
Loss
10
10
10
10
10-5
m1
10a
10a
10a
lteration
Iteration
Figure A8: Training loss and evaluation Le error for learning the solution to 10d Poisson equation over (a) time and (b) steps.
Best run details
The runs shown in Figure A8 correspond to the following hyper-parameters:
SGD: learning rate: 6.550 109 10-3, momentum: 10-1
Adam: learning rate: 1.359 480 10-4
Hessian-free: curvature matrix: GGN, initial damping: 1 10- constant damping: no,
maximum CG iterations: 250
LBFGS: learning rate: 2 10-! history size: 200
KFAC: damping: 1.056 857 10-13 momentum: 7.160 131 10-1, exponential moving
average: 9.622 372 10-1, initialize Kronecker factors identity: yes
KFAC*; damping: 7.978 934 10-",, exponential moving average: 8.950 193 10-1 initialize Kronecker factors to identity: yes
Search space details The runs shown in Figure A8 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 (ll denotes uniform, and Lll log-uniform distribution):
SGD: learning rate: [U([[1 1- 10-3)), momentum: u({0, 3- 10-11011110000 1,-1001,9-10011 10-
Adam: learning rate: (5e(55555 10-
Hessian-free: curvature matrix: U((GGN, Hessian)), initial damping: U({1, 10-1, 10-2,1 10-3,111111 10-4)), constant damping: U4 [({no,, (((500,20))
maximum CG iterations:
LBFGS: learning rate: 10-1,,10-222 2.--1,2.10--1 10-1 1.5 10- 10-2 1.1.2.11110 10--2)) history size: (((00, 75, 100, 125, 150, 175, 200, 200)2255)
KFAC: damping: [ ( 10-14:1- 1- 1- momentum: 4<[(; 9.9 10-1)), exponential
moving average: ((55-10-1;9.3) 10- initialize Kronecker factors to identity: yes
KFAC*: damping: [2/(11100122 1-1-9)),, exponential moving average: (1( 1099-10111 10-1)).
initialize Kronecker factors to identity; yes
23

---

A.5 5/10/100-d Poisson Equations with Bayesian Search
Setup Here, we consider three Poisson equations -Au(x)) = f()) with different right-han hand sides and boundary conditions the unit square x: E [0.1]4:
d with cosine sum right-hand side f(x) X² E cos (TT:), boundary conditions
u(x) = Σi= cos (T5, for æ E @ 1]", and known solution u, (x) = >i= cos (T.)))
We assign each run budget of 0.0
d = 10 with zero right-hand side f(x) 0, harmonic mixed second order polynomial
boundary conditions u(x) = Σ -22- T for æ a[o,, I]4, and known solution u. (x) =
Σ 4/2 Li=1 12- We assign each
budget of s00
d 100 with constant non-zero right-hand side f(x) -2d, square norm boundary
conditions u(x) ||>|| for E 0[0. 1]d, and known solution (x) ||~:| We assign each run a budget of 10 s00
We tune the their associated samplingg 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 A9 visualizes the results.
described in $A.1, as well as the batch sizes No, Nan. and
Best run details
The runs shown in Figure A9 correspond to the following hyper-parametent
5d Poisson equation, 5 + 256 + 256 + 128 + 128 + MLP with D = 116 865
SGD: learning rate: 2.686 653 10-4, momentum: 9.878 243 10-1, Ng: 606, Nao:
2 001, batch sampling firqueecy: 1 170
Adam: learning rate: 6.111 767 10-5, -5- Nn: 534, Nan: 1 1021, batch sampling fre-
quency: 8 220
Hessian-free: curvature matrix: GGN, initial damping: 1.541 954 10", constant damping: no, maximum CG iterations: 358, Nq: 1 084, Nag: 3 3837, batch sampling
frequency: 230
LBFGS: learning rate: 1.749 124 10-1, history size: 339, No: 5 391,, Nag: 4668,
batch sampling frequency: 155
KFAC:
damping: 4.251 10-110 momentum: 9. 986 10- exponential moving average: 9.737 093 10-1, initialize Kronecker factors to identity: yes, Ny: 490,
Nov: 2 708, batch sampling frequency: 2339
KFAC*; damping: 2.240 865 10 12, exponential moving average: 8.522 194 . 10-1,
10-12
initialize Kronecker factors to identity: yes, Nq: 149, Nan: 3801, batch sampling frequency: 1393
10d Poisson equation, - 256 >> 256 + 128 + 128 - MLP with 118 145
SGD: learning rate: 5.805 516 10-2, momentum: 9.715
10-1, 10-1,.00 Ng: 537, Nag:
1 173, batch sampling frequency: 1 083
Adam: learning rate: 1.337 679 10-4, No: 115, No0: 1 960,, batch sampling fre-
quency: 9975
Hessian-free: curvature matrix: GGN, initial damping: 8.963 629 10-1, constant damping: no, maximum CG iterations: 143, Ng: 3 736, Nag: 961, batch sampling
frequency: 3
LBFGS: learning rate: 1.695 334 10-1, history size: 338, No: 342, Nan: 765, batch
sampling frequency: 845
KFAC: damping: 6.575 415- 10-4, momentum: 9.772 500 10-1, exponential moving
average: 2.745 481 10-1, initialize Kronecker factors to identity: yes, Ng: 1884,
Nog: 2 258, batch sampling frequency: 455
KFAC*; damping: 7.530 350 10-12 exponential moving average: 9.648 138 10-1, initialize Kronecker factors to identity: yes, Nq: 1 090, Non: 930, batch sampling frequency: 2 454
100d Poisson equation, 100 + 768 + 768 512 512 + 1 MLP with D = 1 325 057
24

---

(a)
Poisson, 116,865
10d Poisson, 118,145
100d Poisson, 1,325,057
10a
10-1
10"
L error
10-1
1
10
10
10-2
10
10"2
10a
10
Loss
10
10-6
10
10
10-5
10a
10a
100
10
10a
10a
10
10
10
Time [s]
Time [s]
Time [8]
(b)
10a
109
Li eror
10
10-1
10-3
10
10-2
10
10-2
10a
10
Loss
10-1
10
10-5
10"
1075
4/
10
10a
10
10
10a
10a
10
10a
Iteration
Iteration
Iteration
SGD
Adam
Hessiam-frer
LBFGS
KFAC
KFAC
Figure A9;
Training loss and evaluation L2 error for learning the solution to high-dimensional Poisson equations over (a) time and (b) steps using Bayesian search.
SGD: learning rate: 1.450 764 10-3, momentum: 9.747 671 10-1, No: 177, Nan: 2 422, batch sampling frequency: 519
Adam: learning rate: 7.894 685 10-5, Np: 100, Nag: 601, batch sampling frequency:
19
Hessian-free: curvature matrix: GGN, initial damping: 7.705 318 10-4, constant damping: no, maximum CG iterations: 263, Ng: 108, Nag: 2372, batch sampling frequency: 55
LBFGS: learning rate: 1.797 096 10-1, history size: 112, Ng: 2 115, Nag: 1 1852.
batch sampling frequency: 23
KFAC: damping: 9.724 117 10-3 momentum: 5.015 715 10-1, exponential moving
average: 9.200 952 10-1, initialize Kronecker factors to identity: yes, Nq: 124, Nag: 2332, batch sampling frequency: 322
:F damping: 1.236 763 10-7, exponential moving average: 8.302 663 10-1, initialize Kronecker factors to identity: yes, No: 175, Now: 086, batch sampling frequency: 16
25

---

Search space details The runs shown in Figure A9 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 (4 denotes a uniform, and LU log-uniform distribution):
5d Poisson equation, + 256 - 256 + 128 + 128 + MLP with D 116 865
SGD: learning rate: [u((1 10-6; 1)). momentum: u([0 9.9 10-'l), No:
4({100, 101, 10 (00)), Nan: (1150,55,...... 5000)), batch sampling frequency: U( 1..., 10 1000)
Adam: learning rate: cu([1 10-6;1]) No: u{{100, 101, 10000)), Nan:
U(150, 51,... 51,...,,000)) 5000)), batch sampling frequency: u( 1, (,.,, 10000)) initial damping:
Hessian-free: curvature matrix: u([GGN, Hessian)).
CU((1 10 10-15; 15, 1), constant damping: U((no, yes}), maximum CG iterations:
24[1,2..... 500)), No: U({100, 101, 101..... 10000)) Nan: U((50, 5 5000),
batch sampling frequency: U((0, 10000))
LBFGS: learning rate: [U([1 - 10[[1-10-6 II), history size: u(15,6, 6,..., 500}), No:
U(1100, ..... ....,..0000))........... Nan: U(150, ..... (.4....,,000), batch sampling frequency:
1( (000))
KFAC: damping: C4 . 10-15: 10-21). momentum: 4[00 9.9 10o01)-- expo- nential moving average: (10 9.9 1001--11111 'D), initialize Kronecker factors identity: U( No:
101,..., 100000))))) Non: u(150, 5 5000), batch
sampling frequency: (( (..,, 10000))
KFAC*; damping: CU((L 10 10-15;1 10 (0) exponential moving average: U(099.9 - 10 10-1)), initialize Kronecker factors to identity: [(noo, yes)), No:
U({100, 101, 101..... 10000)), Nan: u(150,51, 5 000)), batch sampling frequency:
4(10, 1, ..... 10000))
10d Poisson equation, 10 - 256 - 256 + 128 + 128 - MLP with D = 118 145
SGD: learning rate: LU((1 10-6; 1)), momentum: u([0; 9.9 10-1)), No: (({100, (0l...,,000)), Nag: ((150,61...., 250...,,500),,, batch sampling frequency: 1, 500)))
Adam: learning rate: Lu((1 - 10-6; 1)), No: 4{{100, 101, 101,...,,50000 5000)), Nan:
U({50, 51,.... 2 2500)), batch sampling frequency: u(f0, I...., 5000))
Hessian-free: curvature matrix: u((GGN, Hessian)). initial damping: [M 10 15, 1)), constant damping: 4((no, yes?), maximum CG iterations: u(11, 2..., 500)). No: 4((100, 10,, 5000)), Nan: U((50, 21, 500)))
batch sampling frequency:
LBFGS: learning rate: CU((1 - 10-6; 1l). history size: u(15, U(15,6....,,000 500}). No:
u{{100, 10.,..., (0..)),))))........ Nan: 4([50,,1, 2 2500)). batch sampling frequency:
1, 000))
momentum: U[[0 9.9 10-1)), expo- nential moving average: ((0 9.9 10-1)), initialize Kronecker factors to identity: u[[no, )). Nn: 2({1100, 101, 101,..,, 5000), Noo: U((50,
KFAC: damping: [4(1 10-11-001
batch
sampling frequency: 2(00, ((000))..,,0000)
KFAC*: damping: LU((1 10 10-15; 10 10, exponential moving average: 9(099.9- 10 D) initialize Kronecker factors to identity: u('no, yes)). No:
U(1100, 101,..., 5000)), Non: u(150, 5 2500)), batch sampling frequency: u(10,1, 1,.., 5000))
100d Poisson equation, 100 + 768 + 768 + 512 + 512 + MLP with D = 1 325 057
SGD: learning rate: Lu 10-6; 1)). momentum: u([0; 9.9 10-'1), No: U({100, 101, ...., 5 5003), Non:
51,
batch sampling frequency:
u(10, 1,.., 1 (00}})
AAdmm: learning rate: [U([L 10 IJ), No: U{{100, 101,...., 5000)), Nao:
((55,, 51,.... 2 2000), batch sampling frequency: u((0, I,.., 1 000})
- -essian-free: curvature matrix: u((GGN, Hessian)), initial damping:
LU((1 10-15, 15, I)), constant damping: ((([no,yes)), maximum CG iterations: u([1,2,... 500)), No: 4((100, 101,..., 5000)), Nao: (1,5,,,,......2 2 2500))
batch sampling frequency: U((0,1,... 1000})
26

---

LBFGS: learning rate: LU([1 10-6;1]) history size: U(15. 6..... 500}), No:
U{{100, 101..... 101....50000), Nan: 4({50, 51 .........000,, batch sampling frequency:
4({0, 1,....,....0........ 1000))
KFAC: damping: LU((1 100[1-10-11::11 11:1:10:111 10- momentum: u[ 9.9 10 10-1)), expo- nential moving average: [1; - 10- initialize Kronecker factors to identity: U({no. No:
Noo: U((50, 51,.... , 2000), batch
sampling frequency: ((f0,1,...,10000))
KFAC*: damping: Cu((1 1101-10--110 10-15; (0). exponential moving average: 9.9 10-1)) initialize Kronecker factors identity: U( (((no,yet) No:
4({100, 101,..., .....,,0000 5000)), Nao: u(150,51, ,....,2500 200)), batch sampling frequency:
u(10, 1, ,..... 1000))
A.6 PINN Loss for the Heat Equation
Consider the (d+ -)ddimnssio homogeneous heat equation
Öza(t, ab) ut σ)
with spatial coordinates E B S æd and time coordinate ' S IR where T is a time interval and 6 > 0 denotes the heat conductivity. In this case, our neural network processes (d = d + 1)- dimensional vector æ (t, TT 12d and we can re-write the heat equation as
0z, u(x)
Σ A uu 0.
In the following, we consider the unit time interval = [0; 1]. the unit square = [0; 1] and set types of constraints
1/4. There obtain unique solutions: initial conditions and boundary conditions. As our framework for the KFAC approximation assumes only terms in the loss function, combine the contributions from the boundary and initial values into one term.
need to enforce the heat equation in order to
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 uo w(0, .)) = H sin(TTū;) for ø e Ω. For boundary conditions, use gt, Ø) O for (t, σ) T On. The manufactured solution is
2 dt
u. (t, ②) = exp
sin sin(TT(.)) [ī.]).
4
The PINN loss for this problem consists of three terms: PDE term, an initial value condition term, and spatial boundary condition term,
No
L(
1
Σ
0:uo (:ua(2:))
1 - Da up (2x
Nn
Noa
(uo (xr 000 g(r 002)
+
N Non
No
+
(0, xh)) - up (xx)))
No
12ml
with o ~ x Ω and x ~ T x 02, and Φf ~ {0} x Ω. To fit this loss into our framework which assumes two loss terms, each of whose curvature is approximated with Kronecker factor, we combine the initial value and boundary value conditions into single term. Assuming Nan = No Non1/2 without loss of generality, write
No
Noai
L(O) = No.
1
Σ
(x) - Aa, uo u y
Juo (x (xndd) - 1ond 1|2
Neond n
n=l
L,)(0)
Loonn(
27

---

with domain inputs i ~ T x Z and targets y = O, boundary and initial condition targets grond u. ( (x)en)) with initial inputs In goood
aernnd {0} for 1..... Nome/2 and boundary inputs T x 8n for n Nums/2 + 1,..., Nennd- This loss has the same structure as the PINN loss in
Equation (1).
A.7
1+ld Heat Equation
Setup We consider -t--dimensional heat equation 0,u(t, utt x)) 0 with 1/4 on the unit square and unit time interval, x, t [0, 1] x [0, 1]. The equation has zero spatial boundary conditions and the initial values are given by u(0,
sin( for æ: E [0. 1]. We sample single training batch of size No = 900, Nan = 120 (Nen/2 points for the initial value and spatial boundary conditions each) and evaluate the L2 error separate set of 9 000 data points using the known solution u, (f, three MLP architectures of increasing size, each of whose linear layers are Tanh-activated except for the final one: a shallow 2 - 64 + 1 MLP with D 257 trainable parameters, a five layer 873 trainable parameters, and five layer 2 256 - 256 + 128 - 128 + 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 A10 summarizes the results, and Figure All illustrates the learned solutions over training for all optimizers on the shallow ML.
exp( (0sp(= (Tx).. Each run is limited to 000 We compare
2 - 64 64 48 + + 1 MLP with D
Best details
The runs shown in Figure A10 correspond to the following hyper-parameters.
2 + 64 + MLP with D = 257
SGD: learning rate: 1.752 752 10-2, momentum: 9.9 10-1
Adam: learning rate: 8.629 006 10-4
Hessian-free: curvature matrix: GGN, initial damping: 10 10-4 constant damping:
maximum CG iterations: 350
LBFGS: learning rate: 1 10-1, history size: 125
ENGD (full): damping: 1 10-12 exponential moving average: 10-1 initialize
Gramian to identity: no
ENGD (layer-wise): damping: 1 10-10 exponential moving average: 3 10-1,
initialize Gramian to identity: no
KFAC: damping: 1.273 754-10--
momentum: 7.562 617 10-1, exponential moving
average: 3.611 724 10 initialize Kronecker factors to identity: yes
KFAC*: damping: 1.968 427 10-9, exponential moving average: 9. 7.703 638 10-1,
initialize Kronecker factors to identity: yes
2 64
48 48 MLP with D 9873
SGD: learning rate: 9.276 977 10-2, momentum: 9.9 10
Adam: learning rate: 2.551 515 10-3
Hessian-free: curvature matrix: GGN, initial damping: 1 10-3, constant damping:
no, maximum CG iterations: 200
LBFGS: learning rate: 2 10-1, history size: 125
ENGD (full): damping: 1 10-6, exponential moving average: 9 10-1, initialize
Gramian to identity: no
ENGD (layer-wise): damping: 1-10 exponeetial moving average: 6- 6-10-1 initialize
Gramian to identity:
KFAC: damping: 3.169 186- 10010-1 momentum: 7.075 879- 100 exponential moving
average: 8.860 410 10-1 initialize Kronecker factors to identity: yes
KFAC*; damping: 5.035 695 10-14, exponential moving average: 9.815 164 10-1
initialize Kronecker factors to identity: yes
2 + 256 + 256 + 128 + 128 + MLP with D = 116
SGGD learning rate: 5.709 474 10-2,
momentum: 9.9 10
-
Adam: learning rate: 6.716 485 1
28

---

D 257
9,873
116,097
10
10-2
10-2
Lg error
10
10-4
10-6
10
10
10-3
Mm m
10
WOW Mnma
Loss
10-7
10
10-11
10
10-12
109
10
10
10
10a
10
10a
102
10a
(b)
10-1
10
10-2
Lg error
10-3
10
10-4
10-5
10
10
10-7
10-3
10-3
10
Loss
10
10-7
10-5
10-11
10
11
10-12
10a
10a
10'
10a
10
10'
10a
Iteration
Iteration
Iteration
SGD
Adam
Hessian-free
LBFGS
ENGD (full)
ENGD (layer-wise)
KFAC
Figure A10: training loss and evaluation L2 error for learning the solution 1+1 -i-mmennii heat equation over (a) time and (b). each column corresponds to a different neural network.
Hrssian-free: curvature matrix: GGN, initial damping: 10-2 constant damping:
no, maximum CG iterations: 300
LBFGS: learning rate: 2 10-1, history size: 125
KFAC: damping: 2.576 4 13
momentum: 2.043 -95- exponential moving
average: 9.727 829 10-1, initialize Kronecker factors identity: yes
10111 exponential moving average: 9.765 844 10-1, initialize Kronecker factors to identity: yes
KFAC*; damping: 7.343
493 493
Search space details The runs shown in Figure A10 were determined to be the best via search with approximately 50 runs the following search spaces which were obtained by refining initially wider search (l denotes a uniform, and Cll a log uniform distribution):
- 64 + MLP with = 257
- SGGD: learning rate: LM((1 10-3;1 1. 10-1]), momentum: 4((0,3-10-1 10-1,9 10-1 9.9 10-1))
Adam: learning rate: Lu((5 10-1: 1- 10=
29

---

Optimizer
First step
0.1% trained
1 trained
10% trained
True solution
SGD
Adam
LBFGS
Hessian-free
ENGD (full)
ENGD (layer-wise)
KFAC
KFAC*
Figure A11: Visual comparison learned and true solutions while training with different optimizers for the 1+ ld heat equation using a two-layer MLP (corresponding to the curves Figure A10 left). All functions are shown the unit square (z, f) Ω [0; 1J2 and normalized to the unit interval.
Hessian-free: curvature matrix: u((GGN, Hessian)), initial damping: U({100, 1, 10 2. 10 (-]), constant damping: U( (no, yes}), maximum CG iterations: U(150, 250})
LBFGS: learning rate: U(15 10-1,1111111 2. 1 ., 10-1,011100 5. 10-2,0-2222200 - 1 1, 1011-10-))) 10-2)), history size: U((75, 100, 125, 150, 175, 200, 225, 250})
ENGD (full): damping: 4((1 10-6, 10- 10-10 1. 10-12 0}), exponential moving average: (((0,3310-1, 10-1,9 10-9900,199110--11111 (0-1))) initialize Gramian to identity: U((no, yes})
30

---

ENGD (layer-wise): damping: [[{1--00 11..0001111111""".1 1-10
1-100
11-100...1101"" 1-100-11 0}), exponential
19.9 moving average: 3 (- 100--001,,00000 ,6-10-11 ,9-10
9.9
10 10-1))
initialize Gramian to
identity: (((no, yes))
KFAC: damping: [U((1 10-13;
10-2)) momentum: u(15 10 1,9 10-
exponential moving average: 2e( 9.9 10 10- initialize Kronecker factors to identity:
yes
KFAC*: damping: LU((1 10-13; 1. 10-7)), exponential moving average: U(
10 1 10, initialize Kronecker factors to identity; yes
2 + 64 + 48 - 48 + MLP with = 9 873
SGD: learning rate: [u 10 3;1 10 10)' momentum: 3 1..1-10199
10
10 1 - (0--10-)))
Adam: learning rate: [U( 10-4; 1-111-111 10
Hessian-free: curvature matrix: U((GGN, Hessian)), initial damping: U({100, 1,
10-2.1 1. 110 10-9)), constant damping: 4( nno yes))) maximum CG iterations:
4(150, 250})
LBFGS: learning rate: (115
10-2,2- 20-100--- 10-2,
1,1110-))
history size: U((75, 100, 125, 150, 200, 225,
ENGD (full): damping: W((1 1..110000111111 1..100...11000 10-10.10101011 1. 10 10-12.011100000 exponential
moving average: U ((, 10-1,001,,,0 10-1,------1,-111000
.99- 1-9-10-1)) 10, initialize Gramian to
identity: (4 yes})
ENGD (aaee-wise):: [baperiiiisssuupnngg222201111110
(10::n::2(1111111111111 6.110 1.110 1-10-11-- 0}), exponential initialize Gramian to
moving average: 10-0,,31-0-,6-- (10r::/00000001111-0000000000000 10-1)---
identity: ((noo, yes])
KFAC: damping: 11(1110-1111110011 momentum: u1 9.9 10- exponential
moving average: :,. 10- initialize Kronecker factors to identity: yes
KFAC*; damping: L2((1 10-15;1 10-9)), exponential moving average: 2/
10 10-1; 9.9 10 1)- initialize Kronecker factors identity: yes
2 + 256 + 256 + 128 128 MLP with D 116 097
SGD: learning rate: [u 10 3: 10 10-1)) 'J), momentum: u((0, - 1, 1,6.10
10 10-1, 9.9 10-11)
Adam: learning rate: LU((5 10-4; 10-1))
Hrssian-free: curvature matrix: u((GGN, Hessian)), initial damping: U({100, 1, 10 2,1 10-4,00-4010 1,1 10-6)) constant damping: U( [no, yes}), maximum CG iterations:
u(150, 250})
LBFGS: learning rate: W((5 10-1,2 10-11111111- 10--1.111-1--1111111111111111111 10-1,5, 10-2, 2,2 2,- 1-:1--10)))) 10-2)),
history size: U((75, 100, 125, 150, 175, 200,
- KFAC: damping:
momentum: L<{; 9.9 10 10-1)) "1), exponential
moving average: 4([5 100-1:9.-10-11) initialize Kronecker factors identity: yes
KFAC*; damping: [ 10:14; 10- exponential moving average: U(5
10 10-1:9.9 10 'I, initialize Kronecker factors to identity: yes
A.8 4+1d Heat Equation
Setup We consider 4+1-dimensional heat equation d, x) K. ((t
with 1/4
on the four-dimensional unit square and unit time interval, Φ, E [0, 1]a x [0, 1]. The equation has
spatial boundary conditions u(t, = exp( ( Σi sin(2x;) for f, x E [0, 1] x a(o, 1]4 throughout time, and initial value conditions w(0 x) Di=1 sin(2r:) for
[0, 1jt We sample training batches of size N₂ = 3 000, Nan = 500 (Nony/2 points for the initial value and spatial boundary conditions each) and evaluate the Le error on separate set of 30 000 data points using the known solution u. (t) = cxp(-t) Ei= sin(2r;). All optimizers except for KFAC sample a new training batch each iteration. KFAC only samsaesles every 100 iterations because we noticed significant improvement with multiple iterations on fixed batch. 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-parameterrs see $A.10. 10. The results presented in this section are consistent with this additional experiment (compare the rightmost column of Figure A12
31

---

and Figure A14). Each rum is limited to 3000s. We compare three MLP architectures of increasing size, each of whose linear layers are aant-acctvvted except for the final one: a shallow 5 64
MLP with D 449 trainable weights, a five layer 5 64
64
48 + 48
1 MLP with
D 10.065 trainable weights, and five layer + 256 - 256 + 128 + 128 1 MLP with
D = 116 864 trainable weights. For the biggest architecture, full and layer-wise ENGD lead to
out-of -er-orr errors and are thus not tested. Figure A12 visualizes the results.
(a)
D 449
10,065
D 116,865
10
10
10a
L2 error
10
10
10-2
10
10-4
10
10
10-2
Loss
10"5
10
10
107
10-9
10
10
10
10
10
10
10a
(b)
10
10a
10
Lg error
10-3
10-2
10
10
10a
10-1
10
Loss
10
10
10
10-7
10-9
10-8
10a
105
10
10a
10a
Iteration
Iteration
Iteration
SGD
Adam
Hessian-free free
LBFGS
ENGD (full)
ENGD (layer-wise)
KFAC - KFAC*
Training loss and evaluation L2 for learning the solution to 4+1-d heat equation over (a) time and (b) steps. Columns are different neural networks.
Figure A12:
Search space details The runs shown in Figure A12 were determined to be the best via search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search (u denotes uniform, and Lll a log-uniform distribution):
64 MLP with D 449
SGD: learning rate: 7.737 742 10-3, momentum: 9 -0
Adam: learning rate: 3.708 460 10-3
Hessian-free: curvature matrix: GGN, initial damping: 2 10-1, constant damping:
no, maximum CG iterations: 300
LBFGS:: learning rate: 2- 10-1 history size: 175
ENGD (full): damping: 10-10, exponential moving average: 10-1 initialize
Gramian to identity: yes
32

---

ENGD (layer-wise): damping: 1 10 10-0, exponential moving average: 0, initialize
Gramian to identity: yes
KFAC: damping: 1.000 288 10 momentum: 9.474 108 100 exponential moving
average: 7.783519 10-1, initialize Kronecker factors to identity: yes
KFAC*; damping: 2.965 060 10-$, exponential moving average: 9.574 717 10-1,
initialize Kronecker factors to identity: yes
5 + 64 + + 48 + MLP with 10065
SGD: learning rate: 9.357 973 10-3, momentum: 9 10
Adam: learning rate: 7.801 748 10-4
Hessian-free: curvature matrix: GGN, initial damping: 5 10 10-3, constant damping:
no, maximum CG iterations: 400
LBFGS: learning rate: 1 10-1, history size: 225
ENGD (full): damping: 1 10-$, exponential moving average: 0, initialize Gramian
to identity: yes
ENGD (layer-wise): damping: 1-10-6, exponential moving average: 3- 3-10-1, initialize
Gramian to identity: no
KFAC: damping:
10-11 momentum: 7.660 303- 10-1, exponential moving
average: 9.821 414 10-1, initialize Kronecker factors identity: yes
KFAC*; damping: 1.955 740 10-10, exponential moving average: 9.821 778 10-1,
initialize Kronecker factors to identity: yes
5 + 256 > 256 - 128 + 128 > MLP with D = 116 865
SGD: learning rate: 7.192 473 10-3, momentum: 9 10-1
Adam: learning rate: 5.266 284 10-4
Hessian-free: curvature matrix: GGN, initial damping: 2 10-3, constant damping:
no, maximum CG iterations: 250
LBFGS: learning rate: 2 10-1 history size: 200
KFAC: damping: 8.581 322-10-11 momentum: 8.501 747 exponential moving
average: 9.803 115 10-1, initialize Kronecker factors to identity: yes
KFAC*; damping: 3.405 440 10-14, exponential moving average: 8.445 471 10-1, initialize Kronecker factors to identity: yes
Search space details The runs shown in Figure A12 were determined to be the best via search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search (U denotes a uniform, and Lll a log-uniform distribution):
* 5 + 64 MLP with D = 449
SGD: learning rate: [U[(1 - 10-3; 10 10-2)), momentum: 3( 10-1,9 10-1))
Adam: learning rate: [5 10-4; 1- 1-100-111
Hessian-free: curvature matrix: U((GGN, Hessian)), initial damping: U({1, 1 10 10-1,1 1--'.-1-01 -, 1.1-10 3,0 100 10-41)) +1), constant damping: 2 U((no, yes)),, maximum CG
iterations: 2e((50, 250))
LBFGG:: learning rate: 4({5 1 2. 10-1,.111111111 1000-1,11--1,-11111
1 20- 2. 10-2,11111111 10 10-2)),
history size: (({50, 75, 100, 125, 150, 175,200,
ENGD (full): damping:
10-10 10-11 1 10-12, 0}), exponential moving average: (((0,331001, 100 to identity: Z1 U((no, yes})
10-1 initialize Gramian
ENGD (layer-wise): damping: (((1- exponential moving average:
10 4.1 10 10 1-00-)), 10-1)). initialize
10 .9.9
Gramian to identity: (((no, yes))
KFAC: damping: 12
1-
momentum: U([0 10--10 exponential
moving average: 24[ 10 100--000199990001 10- initialize Kronecker factors identity: yes
:F damping: LU((1 10-13; 1- 10-7)), exponential moving average: 24(
10 10-1, 9.9
10-1l), initialize Kronecker factors identity: yes
33

---

5 + 64 + 64 48 - 48 - MLP with D = 10063
SGD: learning rate: [u 10 3,1 10 10))),, momentum: 4((0,3 1,6.10"1 10- 10-1)
Adam: learning rate: LM((5 10-1; 1- 1010-]]
Hessian-free: curvature matrix: (((GGN,,essian) initial damping: U({1,
10-1.1 10-200-221
10-3, 10- 10-111110-1) constant damping: yes)), maximum CG
iterations: (4((50,250) 250))
11111111.50 LBFGS:: learning rate: 4115 2. 10.1--111.11155
10
10 10-2)),
history size: (({50, 75, 100, 125, 175, 200, 220))
ENGD (full): damping:
1011-10-,111100
10 1 10-11 1- 10-12,00 ,0)),
exponential moving average: U(10,3- to identity: [[(no.
10111011--1 ()) initialize Gramian
ENGD (layer-wise): damping: 1((1-1002
10 -0-11:1:00000 10 10-1--101111111-1 ")).
exponential moving average: -
10 9..
10 10-1)), initialize
Gramian to identity: U(nno, yes))
KFAC: damping: [ 1-11---00900 $1), momentum: 4{ 909 exponential
moving average: 4( 10 10110111,0000 19.9 10-1)) initialize Kronecker factors identity: yes
KFAC*: damping: [ (0-14;1 10-$1), exponential moving average: 4[55
10 9.9 9.9 10 10, initialize Kronecker factors to identity: yes
5 + 256 + 256 - 128 + 128 - MLP with D = 116 865
SGD: learning rate: [u((1 10-3; 10=1))
31 1 10 (1), momentum: 4({0, 3 10- 1.00,66.10 10-1,9-
Adam: learning rate: [u((5 10 10-4; 1-11-01000 1-10011)
Hessian-free; curvature matrix: u( Hessian))) initial damping: U([1, 10 10-1,1 11012.1-2- 1- 10 111-100-1111 10
10 *)), constant damping: U([no, yes)), maximum CG
iterations: (/50 250))
11-1,6 LBFGS: learning rate: 2/(55 10--,,211-1 1.-...1111....1011111111111
5. 10-1,,
10 2, 2. 1:-
10-2 2, 10 (0)))),
history size: 4((50, 75, 100, 125, 150, 20,,
KFAC: damping: 22((1111011:111000
momentum: 21([0; 9.9 1009--10111 10- exponential
moving average: 2105444055110 1::
initialize Kronecker factors identity: yes
KFAC*; damping: CM((1- - 10014, 10-*)), exponential moving average: 4[55
10-1;999 10-1)), initialize Kronecker factors identity: yes
A.9 Robustness Under Model Initialization for 4+1d Heat Equation
Here we study the robustness of our results from $A.8 for the 4+1d heat equation when initializing the neural network differently. We choose the MLP with middle panel which is bigger than the two-layer toy model, while still allowing to run ENGD. Using the same
10 065 parameters from Figure 2's
we re-run all optimizers with 10 different model initializations. The results are shown in Figure A13. We observe that all optimizers perform similar to Figure 2, except for LBFGS which diverges for some runs.
A.10
4+1d Heat Equation with Bayesian Search
Setup We consider the same heat equation in $A.8 and use the 5 256 + 256 128 128 + 1 MLP with D 116 865, We tune all optimizer hyper-parameters as described in $A. and also tune the batch sizes No, Nao, as well their re-sampling frequencies. Figure A14 summarizes the results.
Best run details
The runs shown Figure A14 correspond the following syper--arameters:
SGD: learning rate: 1.614 965 10-2, momentum: 9.899 167 10-1, 1.,,,N,, No: 527, Nan: 2157, batch sampling frequency: 543
Adam: learning rate: 2.583 569 10-4, No: 472, Noo: 3018, batch sampling frequency:
177
34

---

10
10
10
L2 error
10
SGD
Adam
10-3
Hessian-free
LBFGS
(full)
10-1
FNGD (layer-wise)
ENGD (diagona))
KFAC
KFAC*
10a
10
Time
Figure A13: Best runs from the MLP with 10 065 parameters on the 4+ ld heat equation from Figure middle repeated over 10 different model initializations. All optimizers perform similarly, except for LBFGS which diverges for some runs.
(a)
10
10
10 2
La error
SGD
Adam
Loss
10 0
10
LBFGS
10-3
KFAC
10-11
KFAC*
10
10a
10a
10a
10
10
10a
10a
Time [s]
Time [s]
(b)
10-1
10
Lo error
Loss
6
10-3
10
10
10a
10a
10a
10a
10
10a
10
10a
10a
10a
10
10a
Iteration
Iteration
Figure A14: Training loss and evaluation L2 error for learning the solution to 441-dimmnsio heat equation over (a) time and (b) using Bayesian search.
Hessian-freet curvature matrix: GGN, initial damping: 5.077 10-1, constant damping:
yes, maximum CG iterations: 163, No: 172, Nan: 637, batch sampling frequency:
5 440
35

---

LBFGS:
learning rate: 1.029 194 10-1, history size: 488, Np: 582, Nop: 2038, batch
sampling frequency: 315
KFAC: damping: 8.435 180 10-14 momentum: 9.718645 10-1, exponential moving
average: 9.800 744 10-1, initialize Kronecker factors to identity: yes, No: 2 222,, Nan: 2 663, batch sampling frequeney: 7916
KFAC*; damping: 8.837 871 10-15, exponential moving average: 9.887 596 10-1, initialize Kronecker factors to identity: yes, Ng: 2633, Nag: 2873, batch sampling frequency: 647
Search space details The runs shown in Figure A14 were determined to be the best via a Bayesian search the following search spaces which each optimizer given approximately the same total computational time (l denotes a uniform, and Cll a log- uniform distribution):
SGD: learning rate: cu( 10 10-6;1] 1)), momentum: u( 9.9 10-2)), No: U({100, 10.,..., 10000)), Nea: U(150, (((50,51,,.,,500)), 5000)), batch sampling frequency: U((0, 1, 10000))
Adam: learning rate: LU((1 10 1J), Ng: u({100, 101..... 10000)), Naa:
U([50,51, 51, 5000)), batch sampling frequency: U({0, 1, ((0(0,))).,,,0000)
Hessian-free: curvature matrix: u({GGN, Hessian)), initial damping: [U 10-15; 1]). constant damping: yes)), maximum CG iterations: 2,(1,2,... 500)), No: U({100, 101, 10000)), Nao: U((50, 51, 51,..., 5000), batch sampling frequency:
U((0, 1..., 10 100))
LBFGS: learning rate: [u([1 10 1;). history size: u(15, (1,,,,,..,,,,,,,, No: u({100,, 101, ...., 10000)). Nan: U({50, 51,..., 5000)), batch sampling frequency: U({0, 10 10000))
10-21), momentum: 4([0; 9.9 10-1)), exponential moving average: (0; 10-1)), initialize Kronecker factors to identity: 2/( (((no, yes)), Nq: (1000 101,
KFAC: damping: LU(( 100115:1 1-
Nan: 5(550,51,.... 000)), batch sampling frequency:
11),- exponential moving average: U( 10 10111), initialize Kronecker factors to identity: U( ye))), Ng: U({100, 101, ..., 10 000)). Naa: ((500, 51,... ....... 5000)), batch sampling frequency: 4((0, (0.,,,,..... 1000))
KFAC*: damping:
A.11 9+1-d Logarithmic Fokker-Planck Equation with Random Search
For a given drift μ [o, 1] x Rd Re and diffusivity a [0, 1] - Rdxd the Fokker-Plancl equation with initial probability density Po is given by
O.p + div(up) Tr(ao v =0, p00) po:
which is posed on [0, 1] x Rd. Note that p(t, -) is a probability density on Rd for all t E [0, 1]. We transform the above equation into logarithmic space via = log(p). Then & solves
1
0.g div(n)) Vq p -- Vall ir(aa TV2g) =0 <(0) log Po-
For the concrete example of the main text, we set ju(t x:) - - and o /2I E Rdxd, We consider a 9+1 dimensional Bokker-PPannee equation in logarithmic space and replace the unbounded domain by [0, 1] x [-5,5] 5] Precisely, we aim to solve the equation
0,g(t, x))
(t,, x) x [t
Ag x)) =0, ((0) = log(p" (0)),
2
2
where d = 9, E [0. and x E [-5.5] 5]. The solution q" log(p") is given as p² x)) ~
NO exp(-1)1 + (1 exp( -1)22)) The PINN loss includes the PDE residual and the initial
conditions. We model the solution with medium sized tanh-activated MIP with D = 118 145 and the layer structure - 256 - 256 - 128 - 128 - and use batch sizes of No = 3 300, Nag 000. Each run is assigned budget of 6 000 Figure A15 visualizes the results.
36

---

(a)
10a
10a
L2 error
10
Adam
Loss
10
Hossian-free
LBPGS
10-4
KFAC
10
KFACO
10
10
10a
10
10a
10
10a
Time [s]
Time
(b)
10a
10a
La error
LOss
100
10
10
10
mm
10
10a
10a
10a
10
Iteration
Iteration
Figure A15: Training loss and evaluation L2 eπrr for learning the solution to (9+ log Fokker- Planck equation over (a) time and (b) steps.
Search space details The runs shown in Figure A15 were determined to be the best via random search on the following search spaces which each optimizer given approximately the same total computational time (ll denotes uniform, and LUl log-uniform distribution):
SGD: learning rate: [L/([1- 10:3; 3; 10 10) momentum: U({0, 10-1 10000---01..110---0000 ,6-10-1,9-10-1
Adam: learning rate: [Z(([5e-55; 10
Hessian-free: curvature matrix: U((GGN, Hessian)), }), initial damping: U({1, 10-1, 10-2,1 10-3, 1- 10-4)), constant damping: U( yes),, maximum CG
iterations:
(({50, 250))
LBFGS: learning rate: 4((5 10-1, 2.-11..1111111 1. 10-1 10- 2,-2,2210-2, 1.10- size: u({50, 75, 100, 125, 150, 275, 200,225))
10003 10-2)), history
KFAC: damping: [T(( 10.11 10-51), momentum: U([0; 9.9 - 1)) exponential
moving average: ((9999 10 10-1,9.999 10 11- initialize Kronecker factors identity: yes
K damping: [u([1 10-11; 10-51), exponential moving average: U([9.99
10-1; 1001:9.000 9.999 10-11-11-0 "I), initialize Kronecker factors identity: yes
We found that KFAC* requires very large exponential moving averages to work well.
37

---

Update Kronecker factors of the interior loss B Pseudo-Code: KFAC for the Poisson Equation
Algorithm 1 KFAC for the Poisson equation.
Require:
MLP with parameters 0u = (6) ..... 0 (vec W ...., vec W(Z)),
interior data ((wn,ya) )) No
boundary data {(r),, yh h
exponential moving average 3. momentum je, Damping A, number of steps T
Initialization
for
1,....
L do
A for
1 A
052
(1)
B
or
Initialize Kronecker factors
for f O,.... T do
1) Compute the interior loss and update its approximate curvature
Forward Laplacian wit intermediates Compute layer output gradients (1) -000//z!!! with autodiff in one backward pass (1)
( _ Aun) +- Auo, (x)))
1,..., No
sg..s........................... .................... ....)
grad (At,, (2 2(!2))
1, .,..., No. ,.....S S
d +
for all I..... L do
A 3. + (1 N),,
N,S ENO, - 2.:-1)T
B + B 11)
+ (1 - () No En In.s
(TT
end for
Lo (0.) + 2ND E (Aun - yn
> Compute interior loss
2) Compute the boundary loss and update its approximate curvature
(z) (0) (1.) un) <- No, (TFh.)
1,..., Nan
Forward pass with intermediates
Compute layer output gradients (1) = 0un /2 ( with autodiff in one backward pass
(g (1) 0g..") (1)) <- grad(u,, un. (zl
1,
Non
for all ( 1..... L do
Update Kronecker factors of the boundary loss
A + AA (1)
8n + (1 B) Noo E n= 20 (1-1) 21 1-1
B + BB (1)
o0 + (1 () Noa Σ En
Non (0) .T
ga
end for
Lag (0t) + NNo ΣNen (wn y
b Compute boundary loss
3) Update the preconditioner (use inverse of Kronecker sum trick)
1....... ......LL L do
( (A + AA) & ( Ar) + (A en
XI) (B An)
end for
4) Compute the gradient using autodiff, precondition the gradient
(g (1) ...... g (b) +- grad La (0,) Len(O.), (0!)) ....,.? ((
b Gradient with autodiff
for all = 1 L do
& Precondition gradient
A: C g (
Proposed update direction
8:0) () 1-1
if Oelse A
Add momentum from previous update
end for
5) Given the direction 8())
(()) choose learning rate a by line search & update
for
1, L do
Parameter update
8 ad (2)
0 0 + 5 (
end for
end
return Trained parameters 0r
38

---

C Taylor-Mode Automatic Differentiation & Forward Laplacian
PINN losses involve differrnniall operators of the neural network, for instance the Laplacian. Recently, Li et al. [29] proposed a new computational framework called forward Laplacian to evaluate the Laplacian and the neural network's prediction in one forward traversal. To establish Kronecker- factorized approximation of the Gramian, which consists of the Laplacian's 's gradient, we need to know how weight matrix enters its computation. Here, we describe how the weight matrix of linear layer inside a feed-forward net enters the Laplacian's 's computation when using the forward Laplacian framework. We start by connecting the forward Laplacian framework Taylor-mode automatic differentiation [18, 3J, both to make the presentation self-contained and to explicitly point out this connection which we believe has not been done previously.
C.1 Taylor-Mode Automatic Differentiation
The idea of Taylor-mode is forward proppgate Taylor coefficients, i.e. directional derivatives, through the computation graph. We provide brief summary based on its description in [3].
Taylor series and directional derivatives Consider function f Rd R and its K-th order Taylor expansion at a point æ E Rd along direction av E Rd with a E IR,
0f(z)
0 f(x))
f(a) = f(x + ov) = f(x) +
21 v T
v
Ox
Oxx
a
(89 f(x)
+
Σ
Dij U/2 Vig
3!
Ox
in f2is
11,22,85
k
Ok 8kK(()) f(x)
+
K!
Ox
Vig i,,....
f
We can unify this expression by introducing the K-th order directional derivative of f at æ along v, kk f(x)
[KK
f(x)
U [
2
Σ kE....,ix
Ox
Vi
Vix
K
times
This simplifies the uni uni-directional Taylor expansion to
f(a) = f(x + av) = f(x) + adf(x)lu] + a²
a² 2! f()
3!
a gk f(x)(u....,
+..
KT
K
K
g f(x) k!
=: Σ a
where we have used the notation @ to indicate l copics of v, and introduced the k-th order
Taylor coefficient wf E IR of f. This generalizes to vector-valued functions: If f's output was vector-valued, say f(x) IR, we would have Taylor-expanded each component individually and grouped coefficients of same order into vectors wa E TR. such that [we],, is the k-th order Taylor coefficient of the ith component of f.
A note on generality: In this introduction to Taylor-mode, we limit the discussion to the compu- tation of higher-order derivatives along a single direction v. i.e. k f(=))....., v). This is limited though, e.g. if we set = then we can compute 82 fx))( v] v (0¹ f())/a2))) v. We can set e; to the i-th standard basis vector to compute the i-th diagonal element of the Hessian. But we cannot evaluate off off-diagonal elements, as this would require multi- directional derivatives, like 0² f(x) ee, ejzi). A more general description of mode for multi-directional Taylor series along M directions, f(a),.... a.M) = [a + a1 + + aMuu), which require more general vx]
directional derivatives 9 [0 ....... ux] (each vector can be different) are discussed in [25]. We will use this formulation later to generalize the forward Laplacian scheme to more general weighted sums of second-order derivatives in SC.3.
39

---

Composition rule Next, we consider the case where = h is composition of two functions. Starting from the Taylor coefficients w².... w of im(a) = h( + av), the Taylor coefficients wy...., wi of f(a) f(æ av) follow from Faå di Bruno's formula [18, 3]:
w
Σ n!.... ng! glw)) [0 w
(C20)
(part(()
In the above, part(k) is the set of all integer partitionings of k; set of sets. |o|| denotes the length of set E part(k), n; is the count of integer i in o, and w/ h(ar).
Second-order Taylor-mode Our goal is the computation of second-order derivatives of f w.r.t. æ. So let's 's work out Equation (C20) up to order K = 2. The zeroth and first order are simply the forward pass and the nowward-mode gradient chain rule. For the second-order term, we need the integer partitioning of 2, given by part(2) = {{1, 1). (2)}. This results in
w g(wfi).
(C21a)
w Ög( (wn) [w"),,
(C21b)
1
w マ 0d (w))) jws w + 0g(w))) [w]]-
(C21c)
We can also express w/ wa in terms of Jacobian- and Hessian-vector products of g,
w (Jw: wt,
(C22a)
w/ T 0 (stwi))) Grog
(Jmng((u))) wh
2
(C22b)
2
T 0² [o(wn))]
)
wi
w
Note that first-order Taylor-mode (Equation (C22a)) corresponds to the standard forward-mode autodiff which pushes forward error signals through Jacobian-vector products.
C.2 Forward Laplacian
Our goal to compute the Laplacian of f :R R (in practise, c 1),
[lf())) eil
Aaf(x)= = Σ
Σ w E RR
(C23)
@lf(())l(ee
e.])
where e; the i-th standard basis vector, [f(z)], is the j-th component of f(x). and we have introduced the second-order Taylor coefficients wh,, of f along e,, The Laplacian requires computing, then summing, the second-order Taylor coefficients of d d'yyloor approximations (f(x + e;) h=......
Naive approach We can use Taylor- mode differentiation compute all these components one forward traversal. Adding the extra loop over the Taylor expansions we want to compute in parallel, we obtain the following scheme from Equation (C21).
wo g(w!))
(C24a)
(wfs) } i=l....d
{Ög w (un i
(C24b)
) i=l....
0² (ww)) ww widd + @g(ww)))))))) [w
(C24c)
il.....d
Forward Laplacian framework Computing the Laplacian via Equation (C24) first computes, then sums, the diagonal second-order derivatives (w. hes...d Note that we pull the sum
40

---

inside the forward propagation scheme, specifically Equation (C24c), and push-forward the summed second-order coefficients. This simplifies Equation (C24)
w
g(w())
(C25a)
f
{pg ( [ / i=....d
(C25b)
l=.....d
d
2,1
2
8 gggggn ((w()) juf wfil)) + Og(wn))
Σ
226
(C25c)
Equation (C25) is the forward Laplacian framework from Li et al. [29] for computing the Laplacian of neural network. Here, we have derived from Taylor-mode automatic differentiation. Note that Equation (C25) requires less computations and memory than Equation (C24) because we can pull the summation from the Laplacian into the forward propagation scheme.
C.2.1 Forward Laplacian for Elementwise Activation Layers
We now describe Equation (C25) for the case where g : IRC RR acts element-wise via R - IR. We will write ()),, o' a
to indicate element-wise application of a, its first derivative a', and second derivative o" to all elements of Further, let @ denote element-wise multiplication, and 62 element-wise squaring. With that, we can write the Jacobian as (x)=))(<)) = diag(r(h(<>) where diag( embeds a vector into the diagonal of a matrix. The Hessian of component is 8 (I((h())))//an (2) = (e) Inserting Equation (C22) into Equation (C25) and using the Jacobian and Hessian expressions of the element-wise activation function yields the following forward Laplacian forward propagation:
w! ow")),
(C26a)
1 J
(o'(w)) @ w
(C26b)
Σ W2,1 2 (wo6)) O
1
(w) 02
+ +w(()) o
Σ w
(C26c)
C.2.2 Forward Laplacian for Linear Layers
Now, let g : Pr. + R Dae be a linear layer with weight matrix W R Du and bias vector b E lR Dwi, Its Jacobian is JA(w) (W((x) + b) W and the second-order derivative is zero. Hence, Equation (C25) for linear layers becomes
w
W w/ - b. { ( w/, (=
(C27a)
{ (wff)
(C27b)
l....d
w W
Σ w
(C27c)
We summarize Equation (C27) in single equation by grouping all quantities that multiplied by W into a single matrix, and appending a single row of ones or zeros to account for the bias:
wo wf.1 wh. Ehrw! (W
b)
w
Σi w/
=T Per x(2+2)
T PP, 1)x(242)
or, in compact form,
T WTh
(C28)
Equation (C28) shows that the weight matrix ( = (w()) (()) of linear layer fO) inside a neural network f(T) o... f()) applied to a matrix ( E R Pax (d+2) during the computation of the net's prediction and Laplacian via the forward Laplacian framework and yields another matrix
T(() E R Dau (d+2).
41

---

C.3 Generalization of the Forward Laplacian to Weighted Sums of Second Derivatives
The Laplacian is of the form Azf = E, d f(z)]e, e,] and we previously described the forward Laplacian framework of Li et al. [29] as a consequence of pulling the summation into Taylor-mode's forward propagation. Here, we derive the forward propagation to more general operators of the form En, or [m)) ej], which contain the Laplacian for Ci.) وū.j.
As mentioned $C.1, this requires generalization Taylor-mode which computes derivatives of the form 0 (o=)]o.... ...... v). where the directions must be identical. We start with the formulation in [25] which expresses the K-th multi culti-dirrectioall derivative of a function f o through the composites' derivatives (all functions be vector-to-vectee
0k f( ukl
Σ
al (l(x))
(()))))))(()) [i)
(C29)
(Epart(f1.....K])
Here, part(f1...... K}) denotes the set of all set partitions of {1,.... K (ar is a set of sets). E.g.,
part((1)) (((1))),
([1,)[ ([(1,2)), (0), (2))),
part((1,2, = =(f(1,23))) ((1), (2.3)). ((1,2), (3)). (f1,3) (2)) ((1), (2), (3))).
To make this more concrete, let's consider Equation (C29) for first- and second-order derivatives,
[f (g(h(()))0h(<))||],,
(C30a)
8 f(=)(v), u2] 0g²² (h(z)) [m)] [b(]] + 0g(h(=))))h v2]l
(C30b)
From Equation (C30), we can see that if we want to compute a weighted sum of second-order derivatives Eis C. 02 f()) Uyl, we can pull the sum inside the second equation,
Σ a.j0 f()) (2 2y} = Σ (O(())))0 (x(=)(,]]
(C31)
=g(h(()))
Σ C my]
Hence, we can propagate the collapsed second-order derivatives, together with all first-order deriva- tives along v1, ....... The only difference to the forward Laplacian how second-order effects of an operation are incorporated (first term in Equation (C31)).
We now specify Equations (C29) and (C31) for linear layers and element-wise activation functions.
For linear layer g h(x)) Wh(x) b. we have gm>> (h(())) eml 0, and thus
(f())(o) W Wäh
(C32a)
8 fxx) l2a W / (x))(u1, 02],
(C32b)
a f(s)(v)..... ukl ] = W [(=))(v)..... uk].
(C32c)
The last equation is because only the partition {1, (1...., K) contributes Equation (C29).
For elementwise activations h(xx) + [h(()) with a R + R applied component wise, we have (0(''(((()))))))))") Do,,.... and multi-directtonal
the structured derivative tensor [0Th ((((()))),,.......
derivative kK g(h(())[[1,,..... uk| 2 (x)) oor C Ak. Equation (C30) becomes
[f(x)(v) = (1()))) o (p),
(C33a)
fl v2] = ==(((x)) @ [h( @ [b(2)(vz) + o' d'(h))) @ Ph(z)(u,, v2).
(C33b)
As shown in Equation (C30b), for both Equations (C32) and (C33), we can pull the summation inside the propagation scheme. Specifically, compute Σij d² 82 ((x) le,, ez], we have for linear layers
f(x) = g(h(z))
(C34a)
0f(z)]e.} w wah(x)le))
i I.
(C34b)
Oa w Σ ci Ph()eeeee
(C34c)
42

---

and for activation layers
f(x) = d(((z)),
(C34d)
[f())(e) = = d' (h(x)) & [e))) i=
(C34e)
a ex) Σ Cayo Eayo"(((x)) (h(x)) o ah(x)ee o @h(x)le,}
%d
(C34f)
+ ( O
cr Ph())ee
(the summed second-order derivatives that porpard-propagated are highlighted). This propagation reduces back to the forward Laplacian Equations (C26) and (C27) when we set Ci.j 6i.- In contrast other attempts compute such weighted sum of second-order derivatives by reducing it to (multiple) partial standard forward Laplacians [30], we do not need to diagonalize the coefficient matrix and can compute the linear operator in one forward propagation.
C.4
Comparison of Forward Laplacian and Autodiff Laplacian
Setup
We compare the efficiency of the forward Laplacian, that we use in all our experiments, to off-the shelve solution. We consider two Laplacian implementations:
1. Autodiff Laplacian. Computes the Laplacian with PyTorch's automatic differentiation (functorch) by computing the batched Hessian trace (via torch. func. hessian and torch. func. vmap). This is the standard approach in many PINN implementations.
2. Forward Laplacian. Computes the Laplacian via the forward Laplacian framework. We used this approach for all PDEs and optimizers, except ENGD, presented in the experiments.
We use the biggest network from our experiments (the Dn + 768 + 768 + 512 + 512 + MLP with tanh- activatiiss from Figure 3), then measure time and peak memory of computing the net's Laplacian on a mini-batch of size N = 1024 with varying values of DR. To reduce measurement five independent Python sessions and report the smallest value (using
noise, repeat each the same GPU as all other experiments, an NVIDIA RTX 6000 with 24 GiB memory).
Results
The following tables compare run time and peak memory between the two approaches:
Autodiff Laplacian [s]
Forward
Autodiff
Forward
Da
Dg
Laplacian [s]
Laplacian [GiB]
Laplacian [GiB]
0.051 (1.6x)
0.033 (1.0x)
0.21 (0.96x)
0.22 (1.0x)
10
0,20 (2.0x)
0.10 (1 100)
10
0.98 (1.6x)
0.61 (1.0x)
100
1.7 (2.0x)
0.84 (1.0x)
100
8.8 (1.9x)
4.6 (1.0x)
We observe that the forward Laplacian is roughly twice as fast as functorch Laplacian, and that it uses significantly less memory for large input dimensions, up to only half when Dn 100. We visualized both tables using more values for Dg. see Figure C16. In the shown regime, we find that observe linear
the MLP's increasing cost in Dn (due to the growing first layer) is negligible scaling in both memory and run time. For extremely large Da. it would eventually become quadratic.
D Backpropagation Perspective of the Laplacian
Here, derive the computation graphs for the Laplacian and its associated Gramian when using reverse-mode AD, aka backpropagation. In contrast to the Taylor-mode perspective, the resulting expressions cannot be interpreted simple weight-sharing. This complicates defining Kronecker- factored approximation for the Gramian without introducing new approximations that are different from Eschenhagen al. [17], rendering the Taylor-mode perspective advantageous.
We start by deriving the Laplacian Au := (T) of a feed-forward NN (see $2.1), assuming
single data point for simplicity (see $D.1) and abbreviating uo The goal is to make the
Laplacian's 's dependence w.r.t.a a weight W in one layer of the network explicit. Then, we can write
43

---

Forward Laplacian functarch Laplacian
Peak memory [GiB]
8
Time [s]
6
0.5
2
0.0
20
80
100
0 20
40
100
Da
Dn
Figure C16: Time (left) and memory (right) required with the forward Laplacian used in our implementation and the functorch implementation.
down the Jacobian Jwi Au (see SD.2) which is required for the Gramian Equation (2) (see $D.3). We do this based on the concept of Hessian backprepagation [10, HBP.J which yields recursion for the Hessian Vzu. The Laplacian follows by taking the trace of the latter. Finally, we use the chain rule express the Laplacian's Jacobian J woo Au terms of W children the compute graph.
D.1 Hessian Backpropagation and Backward Laplacian
Gradient backpropagation describes a recursive procedure to compute gradients by backpropagating a signal via vector-Jacobian products (VJPs). A similar procedure can be derived compute Hessians w.r.t. nodes in graph (z(1) (()).. We call this recursive procedure Hessian backpropagation [10].
Gradient backpropagation As a warm-up, let's recall how compute the gradient Vou (Vomy,.... Voucy u). We start by setting Vg(o) u = V.u = (assuming u is scalar for simplicity). then backpropagate the error via VJPs according to the recursion
(Jzr 21)
V 2t 1)
(D35)
( V (Join
for L,...,
1. This yields the gradients w.r.t. intermediate representations and parameters.
Hessian backpropagation Just like gradient backpropagation, we can derive recursive scheme for the Hessian. Recall the Hessian chain rule
v)( (Jg) 2 (J) E(vgDe Vagk
(D36)
where zi denotes the individual components of g. see [53]. The recursion computing Hessians of w.r.t. intermediate representations and parameters starts by initializing the recursion with V2 u
V u O, and then backpropagating according to (see Dangel et al. [10] for details)
22 211-1; (J
Viig ( (1)
22
(D37)
A(t)
2 0z; (J)))=(0)
V2co (J
Σ
Voico (=()) (2c
+
for J....,, 1. The first term takes the incoming Hessian (w.rtt. a layer's output) and sandwiches between the layer's Jacobian. It can be seen as backpropagating curvature from downstream layers. The second term adds curvature introduced by the current layer. It is only non-zero the layer is nonlinear. For linear layers, convolutional layers, and ReLU layers, it is zero.
Following the Hessian backpropagation procedure of Equation (D37) yields the per-layer parameter and feature Hessians 2 u, V2 o(o) u. In Figure D17 we depict the dependencies of intermediate gradients and Hessians for computing Vzu 2 22.0
44

---

Parameters
00
(3-1)
0
0
Forward
Backward
Vir 24
200
Vitt
Vau
Hess, hackward
AT
V 22
V
Figure D17: Computation graph of a sequential neural network's Laplacian Au when using (Hessian) backpropagation. Arrows indicate dependencies between intermediates. Note that 2( æ. z(4) = u, V. 1, and V O. For the Gramian, we interested in how the neural network parameters enter the Laplacian's 's computation. Each parameter is used three times: during (i) the forward pass, (ii) the backward pass for the gradient, and (iii) the backward pass for the Hessian.
(t depends on (2(0) u. due the recursion in Equation (D35), and on ((--)), ((() due the Jacobian J (3 x in the gradient backpropagation Equation (D35).
v 1y u depends on ( and Vg(o u due to the recursion in Equation (D37), and on z ((-1), g due to the Jacobian 2=(1-1)2()) and Hessian V2 1z(--)[((1)) in the Hessian back-
propagation Equation (D35).
The Laplacian Au follows by taking the trace of V2 from above, and is hence recursively defined. To make these expressions more concrete, we now recap the HBP equations for fully-connected layers and element-win nonlinear activations.
Hessian backpropagation through nonlinear layers We mostly consider nonlinear layers without trainable parameters and consist of a componentwise nonlinearity + o(z) for some a: IR + R. The Jacobian of such nonlinear layer is given by 2= = diag(o' (z (i-1))) and the Hessian terms given by V - [= 1 = o" (z) (f-1) Jere where ex the unit vector along coordinate k. With these two identities we can backpropogate the input Hessian through such layers via
2 2 (ding(o'(>('-)))
D2r0 (diag(o'(x((1))
h0)
(D38)
o" (x
ee [Vag
Hessian backpropagation through a linear layer To de-clutter the dependency graph of Fig- D17, will now consider the dependency of Au the weight of single layer. We assume this layer to be linear layer with parameters ( such that ( vec wee(W())).
W w() (i-1).
(D39)
For this layer, the second terms in Equation (D37) disappears because the local Hessians
that is V2 [0] O and (
0. Also, the Jacobians are Jw (() = 2(i-1) @
and J 2 (+) and hence only depend of the two layer inputs. This simplifies the computation graph. Figure DI8 shows the dependencies of W(() on the Laplacian, highlighting its three direct children,
2 (i-1) )z(-1)== W T
(V
(D40)
2 -(ー-)= -= T (V ( W()
45

---

Figure D18: Direct dependencies of linear layer's weight matrix W in the Laplacian's 's computation graph. There are three direct children: (i) the layer's output from the forward pass, (ii) the Laplacian's 's gradient w.r.t the layer's input from the gradient backpropagation, and (iii) the Lapla- cian's Hessian w.r.t. the layer's input from the Hes- sian backpropagation. The Jacobians J w Au required for the Gramian are the vector-Jacobian products accumulated over those children.
14
V
Au
291
D.2
Parameter Jacobian of the Backward Laplacian
Recall that the entries of the Gramian are composed from parameter derivatives of the imput Laplacian, see Equation (2). We have identified the direct children of W in the Laplacian's s compute graph, Equation (D40). This allows us compute the Jacobian Jw Au by the chain rule, i.c. by accumulating the Jacobians over all direct children,
J Jw Au E Fa--na (Jwwe)) V.Au
VemAn
(D41)
Vp VF Verennaa + (Jwe) V2(--1))
+ (Jwo
Vpz Au. Au.
((-1)
The terms v.Au be computed with gradient backpropagation to the respective intermediates.
D.3 Gramian of the Backward Laplacian
With the Laplacian's Jacobian from Equation (D41), we can now write down the Gramian block of the interior loss (up summation over the data) for ( as
(J win Au) (Jwis Au)
E.,e(f Vo- Va oJwwoo) (V.Au)) (V.Au)
(Jwio
(D42)
(1)
The Gramian consists of nine different terms, see Figure D19 for a visualization which shows not only the diagonal blocks G(O) but also the full Gramian Go which decomposes in the same way. The terms V. Au u are automatically computed when computing the gradient the loss via backpropagation. We will now proceed and simplify the terms by inserting the Jacobians into Equation (D41) and studying the Gramian's 's block diagonal, which is approximated by KFAC, in more detail.
Computing Jwa Let first compute the Jacobians Jw in Equation (D41). The Jacobian of linear layer's 's forward pass is
Jw (Wx) x @ I. The Jacobian from the gradient backpropagation is
(D43a)
Jw (W x) = I xT.
(D43b)
and the Jacobian from the Hessian backpropagation is
Jw ( xW) & W X
(D43c)
where K E Idim (Ziim(Z)xci) (d)mmm((2))))) 2) (denoting Z = W XW) is a permutation matrix that, when multi- plied onto a vector whose basis corresponds to that of the flattened output Z, modifies the order from first-varies-fastes to aast-varies-fastes.. i.e.
K vec( / vecl Z
46

---

Full interior Gramian
1
(forward, forward)
(forward, gradient)
(forward, Hessian)
(gradient, forward)
(gradient, gradient)
(gradient, Hessian)
(Hessian, forward)
(Hessian, gradient)
(Hessian, Hessian)
Figure D19: Contributions Gp. to the Laplacian's Gramian Go from different children in the computation graph synthetic toy problem. We 4 + 2 1 sigmoid-activated MLP and 10 randomly generated inputs. The contributions are highlighted as Equation (D42).
Re-introducing the layer indices, the expressions in Equation (D40) become
(w) = TT
J Jwn Vea Du, = 8 V V
(D44)
J Jwi 2 =
(
( 1
K
W (V (V)
We will now use symmetries in the objects used during Hessian backpropagation to simplify this further. At a first glance, it looks like the Gramian consists terms, as there are summands from the Jacobians in Equation (D43). However,
simplify into 9 terms:
47

---

First, zl: u is symmetric, that is
Jw and the transposed Jacobian is
W
u) = I & (WOOT (V [V) + K (I (I (wO) (V))
No (V w + ( (V w()l K
Second, we multiply the transpose Jacobian onto V Au, which inherits symmetry from the
Hessian, VV " Aulja VV 1) AuAul, Due to this symmetry, the action of K (or K does not alter it,
KT (Vrg (
Vpa 211-1) Au. Au.
In other words, it does not matter how flatten (first- or last-varies-fastes)) This simplifies the VJP (last line Equation (D42)) to
I @
(V w) wrll))
V V Au ( 2 ( u) wrio)) K T Vpz 11 Au
x (V w()) V
2(1-1) Au.
We can now write down the simplified Jacobian from Equation (D41), whose self-outer product forms the Gramian block for linear layer's weight matrix.
J Jw Au (2(-11)
Var Au
(1)
+ (I 2 V2(12)
V An
(D45)
(2)
+2
2 21 w
V
21--1 Auu
(3)
where (1) is the contribution from the forward pass, (2) the contribution from the gradient back- propagation, and (3) is the contribution from the Hessian backpropagation. The Jacobians from Equation (D43) allow to express the Gramian terms of cronecker-structuree expressions consisting of 9 terms in total. Figure D19 shows the 6 contributions from different children pairs.
Conclusion One problem of computing the Laplacian and its Jacobian with backpropagation ac- cording to Equation (D45) is that if we write out the Gramian's 's block G Jw (u Au) T we obtain terms of different structure. Defining single Kronecker product approximation would involve introducing new approximations on top of those employed by Eschenhagen et al. [17]. There- fore, the forward Laplacian, or Taylor- perspective we choose in the main text is advantageous as it allows to define KFAC without introducing new approximations.
48

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We list our contributions as bullet points in $1 and provide references to the parts the paper that present them.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See the paragraph dedicated limitations $5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
49

---

Justification:
Whereas we provide a derivation of the proposed Kronecker-factore approxi- mation, our work does not provide any rigorous mathematical statements.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed description of the general experimental protocol in $A.1 and details for each individual experiment in $A, including all hyper-paramete search spaces and paper-pprraeters of the best runs. We also show pseudo-code for our algorithm in $B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
50

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will open-source our KFAC implementations, as well as the code to fully reproduce all experiments and the original data presented in this manuscript.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.ce
cttic//guidee//ooeeeeeeee wtt.:Suwwmmsssoooosooooooo onsiolcc) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. wtcppbblcc/guuiiiceeeeeeeeoppeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We list all details in $A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: While we not show mean and standard deviations over different model initializations for training curves, aimed to thorough tuning protocol to avoid artifacts from insufficient paper--arameter tuning and also conducted experiments with alternative hyper- hyper-parantee parameter search methods (Bayesian versus random/grid) to ensure the consistency of our results. To further probe the robustness of our results, we analyzed their stability w.r.t. random seed in for experiment in $A.9.
Guidelines:
The answer NA means that the paper does not include experiments.
51

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errоrs is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We describe the hardware in $4. The total computation time can be inferred from the description of the tuning protocol in combination with the assigned time budget per run.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't ' make it into the paper).
paper (e.g., preliminary or failed experiments that
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https //tn///nneuupp// httpbi/ce/thcssciddeeeeeeeeeeeee
Answer: [Yes]
Justification: We have read the Code of Ethics and believe that our submission conforms to it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
52

---

Justification: It the goal of this paper to improve the accuracy of neural betwork-basee PDE solvers. As the numerical solution of PDEs is a classic topic applied mathematics and engineering, we believe there is no immediate negative societal impact. Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does release any data models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-us should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
The paper does not use any existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
53

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi werrswiihhhheee thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's ' creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release any new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects.
Guidelines:
54

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
55


Reviews:
[{"review_id": "fjswa13z9o", "reviewer": "Reviewer_TQnM", "summary": "This work considers the problem of optimising partial differential equations (PDE) with neural networks, in particular second-order optimization.\nEven if simple models (multi-layer perceptrons) are used, for which KFAC approximations of the curvature matrix are well-known, it needs to derive new approximations because the loss function used for solving PDEs is different than previous work (square loss or cross entropy), in particular it includes derivatives of the model with respect to the input, which is interpreted as a larger model with weight sharing. Therefore, it uses the previously (but recently) developed technique of KFAC with weight sharing. Experiments show that the new method is successful when applied to a few example PDEs.", "questions": "NA", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and seems correct, both the theory and experiment. It seems to be the first application of KFAC to PDEs and has the potential of a high impact.", "weaknesses": "Even if experiments with a relatively large number of parameters are shown (10^5), the paper does not provide the complexity of the algorithm, in terms of relevant quantities, for example the input dimension, output dimension, batch size, number of parameters, etc. For example, does the method scale poorly with the input dimension? I would expect quadratic scaling with input dimension. I have not read the forward laplacian paper, but I have the feeling that improves scaling by only a constant factor. https://arxiv.org/abs/1206.6464 argued that it may not exist an algorithm for computing the diagonal of the Hessian that is linear in the dimension using automatic differentiation. Also, as with other KFAC approximations, it remains unclear how general is the approach, e.g. how to apply it to higher order PDEs or other neural network models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer TQnM,\n\nwe would like to thank you for the time and effort you put into reviewing our work. In the following, we want to address some of the points that you rightfully raised.\n\n---\n\n> paper does not provide the complexity of the algorithm\n\n> does the method scale poorly with the input dimension? I would expect quadratic scaling with input dimension.\n\nYou are right that we currently do not discuss this; we address this in detail in our global rebuttal and will add a discussion to the main text.\n\nThe quick summary is that empirically we observe linear run time and memory scaling in the input dimension for computing the Laplacian either via `functorch`'s nested first-order autodiff or the forward Laplacian framework. This linear scaling is intuitive from the forward Laplacian perspective, which performs `input_dimension + 2` forward passes of the net (assuming the network's forward pass is independent of the input dimension, which is a good approximation for the MLP we used, but in general will depend on the architecture). In Landau notation, the scaling with the input dimension will depend on the choice of architecture.\n\nMoreover, we observe some practical advantages of the forward Laplacian over the nested first-order autodiff approach:\n- The constants in the scaling of the forward Laplacian are smaller than those in the `functorch` implementation. The forward Laplacian is roughly 2x faster. We think this is because its computation graph is less sequential. Also, we observe that the forward Laplacian uses roughly half the memory. These constants matter in practise.\n- Backpropagation on the forward Laplacian graph (which we need to compute the gradient of the loss function) allows us to populate the K-FAC matrices without much re-computation.\n\nOverall, it seems that the forward Laplacian is the current state-of-the-art method to compute input Laplacians, comes with advantages over nested first-order autodiff, and integrates nicely with our proposed K-FAC approximation.\n\n---\n\n> it remains unclear how general is the approach, e.g. how to apply it to higher order PDEs or other neural network models.\n\nOur approach is general:\n\n- The proposed K-FAC approximation generalizes to arbitrary PDEs, both in the PINN formulation (i.e. least squares of the strong form) and for variational formulations (deep Ritz). We discuss this in Section 3.3, see also equation (14).\n- The Taylor-mode AD needed to assemble the K-FAC matrices works for general neural network architectures. For instance, JAX offers a general purpose Taylor-mode implementation ([`jax.experimental.jet`](https://jax.readthedocs.io/en/latest/jax.experimental.jet.html)) onto which our proposed K-FAC framework could be built on.\n- Like the traditional KFAC, our KFAC approximation is architecture-dependent. However, thanks to the recent formulation of KFAC for linear layers with weight sharing [2], a wide range of architectures is covered, e.g. fully-connected/attention/convolution layers.\n\n---\n\nThanks again for helping us to improve our manuscript! We hope our responses address your comments. Please let us know if you have follow-up questions. We would be happy to discuss.\n\n## References\n\n[2] Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., & Hennig, P. (NeurIPS 2023). Kronecker-factored approximate curvature for modern neural network architectures.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "inx579iUJs", "reviewer": "Reviewer_YHUU", "summary": "This paper generalizes the K-FAC method (which is well-known in optimization for deep learning) to enable it to train PINNs. The key idea is to combine Taylor-mode autodifferentiation with recent work on K-FAC with weight sharing layers. The proposed method is evaluated on several PDEs, where it outperforms first-order methods like SGD and Adam and quasi-Newton methods like L-BFGS, while performing comparably to (matrix-free) ENGD.", "questions": "* Line 64: Could the authors please elaborate on why the preconditioner not having a Laplacian term is beneficial?\n\n* Line 131: Could the authors please elaborate on why the interior Gramian cannot be approximated with the existing K-FAC? Is this due to computational and/or autodiff limitations?\n\n* Is using a general loss $\\ell$ in equation (11) necessary? For PINNs, $\\ell$ is almost always the least-squares loss, but I could also understand if the authors wanted to show that their method works in a more general setting.\n\n* Does the proposed approach generalize to inverse problems involving PINNs? If it does generalize, how does it affect the calculation of the Kronecker-factored approximation?\n\n* Are there operators besides the Laplacian for which the weight sharing can be reduced?", "limitations": "Please see \"Weaknesses\"", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The paper makes a novel contribution by extending K-FAC to PINNs, which had not been previously done in the literature.\n\n* The paper does a good job of showing how the Kronecker-factored preconditioner can be extended to the residual portion of the PINN loss. I really appreciated the example derivations in sections 3.2 and 3.3!\n\n* The paper empirically evaluates their methods against other second-order/quasi-Newton methods used in training PINNs, such as ENGD, matrix-free ENGD, and L-BFGS. The comparison appears to be fair, since all methods receive equal amounts of hyperparameter tuning.", "weaknesses": "* The paper is missing more recent work on difficulties of training PINNs, such as [1, 2].\n\n* K-FAC/K-FAC* do not reach the lowest error obtained by running ENGD with $D = 449$ (Figure 2). Why is this the case?\n\n* I don’t think the PDE settings tested in this paper are particularly challenging. For example, the 2d Poisson equation and (4+1)d heat equation have solutions that have relatively low frequencies, but PINNs are known to struggle more with learning solutions containing high frequencies [3, 4]. How would K-FAC/K-FAC* fair in more challenging settings [5, 6]?\n\n* The approach for computing the Kronecker-factored approximation uses “a larger net with shared weights”. Could this cause out-of-memory issues with larger networks? How much additional memory is required?\n\n* I think notation is being reused too much in equations (5) and (6). For example, I don’t think it’s a good idea to have $z^{(l)} = W^{(l)} z^{(l - 1)}$ and $z^{(l)} = \\sigma(z^{(l - 1)})$.\n\n* I’d recommend defining the symbol $\\odot$ as the Hadamard product before it appears in equation (6).\n\n* Equation below line 168: $N_\\Omega$ in the sum should be $N$.\n\n* I believe the outer product decomposition in line 200 is missing a transpose on the second $l_{n, m}$ in the summation.\n\n* There is no per-step computational complexity provided for the proposed methods.\n\n* No open-source implementation is provided for the proposed methods.\n\nMinor comments:\n\n* The sentence containing equation (3) is a bit confusing. Saying “$u_n = u_\\theta(x_n)$, and it coincides with the classical Gauss-Newton matrix” makes it sound like $u_n$ itself is the classical Gauss-Newton matrix (although I can see that this is not what the authors meant). \n\n* For consistency with existing literature, I would recommend writing “KFAC” as “K-FAC” and “LBFGS” as “L-BFGS”.\n\n[1] “On the Role of Fixed Points of Dynamical Systems in Training Physics-Informed Neural Networks.” Rohrhofer et al. (2023)\n\n[2] “Challenges in Training PINNs: A Loss Landscape Perspective.” Rathore et al. (2024)\n\n[3] “On the eigenvector bias of Fourier feature networks: From regression to solving multi-scale PDEs with physics-informed neural networks.” Wang et al. (2021)\n\n[4] “When and why PINNs fail to train: A neural tangent kernel perspective.” Wang et al. (2022)\n\n[5] “PINNacle: A Comprehensive Benchmark of Physics-Informed Neural Networks for Solving PDEs.” Hao et al. (2023)\n\n[6] “PDEBENCH: An Extensive Benchmark for Scientific Machine Learning.” Takamoto et al. (2022)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer YHUU,\n\nwe would like to thank you for the time and effort you put into reviewing our work. In the following, address some of the points that were rightfully raised by you.\n\n> How much additional memory is required [for our KFAC]?\n> There is no per-step computational complexity provided\n\nIn our general response, we compare memory requirements and computation time of the proposed scheme. We find that the viewpoint of 'a larger net with shared weights' alias the forward Laplacian graph saves both memory and compute time when compared to the standard approach relying on nested backpropagation, as implemented in `functorch`. If there are any further questions, please do not hesitate to contact us.\n\n> No open-source implementation is provided\n\nWe will release the code after acceptance with the camera ready version.\n\n> K-FAC/K-FAC* do not reach the lowest error obtained by running ENGD [...]. Why is this the case?\n\nAs KFAC is designed as an approximation of ENGD we do not expect KFAC to provide results superior to ENGD. Note however that ENGD requires the solution of a system of linear equations in the parameter space of the network. Hence, while we can use it as baseline for a small network, it does not scale to larger nets. Our updated discussion of Figure 2 elaborates on this.\n\n> How would K-FAC/K-FAC* fair in more challenging settings [5, 6]?\n\nWe believe that PINNs should be employed for high-dimensional PDE problems where classical methods are intractable and are currently working on the log-Fokker-Planck equation which is a challenging nonlinear PDE. We aim to present first results in the discussion period and refined results for a camera-ready version.\n\n> The paper is missing more recent work on difficulties of training PINNs, such as [1, 2]\n\nThank you for the pointers to these references, which we have included into our introduction. \n\n**Notation/typos**\n\n> I think notation is being reused too much in equations (5) and (6)\n\nThank you for this feedback. We are currently looking into ways to streamline and improve our notation for the camera-ready version.\n\n> I’d recommend defining the symbol as the Hadamard product before\n\nThanks for the pointer, we have added the definition.\n\n> Equation below line 168\n\nThanks for spotting this, indeed, we have been slightly inconsistent in carrying the subscript $N_\\Omega$ for the number of interior integration points. We have adjusted this and have made our notation consistent.\n\n> the outer product decomposition in line 200 is missing a transpose\n\nThanks for catching this typo, which we have corrected.\n\n**Minor comments:**\n\n> The sentence containing equation (3) is a bit confusing.\n\nThanks for the feedback. To our knowledge the matrix $J_\\theta^\\top u J_\\theta u$ is usually called Gauß-Newton matrix, where $J$ is the Jacobian with respect to the model parameters. This is exactly the form of the matrix defined in (3) if $u = (u_1, \\dots, u_n)^\\top$. Please let us know, whether you have any further questions and we are willing to adjust our wording to reduce ambiguities.\n\n> I would recommend writing “K-FAC” and “L-BFGS”.\n\nThank you for the feedback, of course we are willing to follow the common nomenclature.\n\n**Questions**\n\n> Line 64: Could the authors please elaborate on why the preconditioner not having a Laplacian term is beneficial?\n\nWe are deeply convinced that when optimizing a PINN, preconditioners with a PDE term, e.g., a Laplacian, are much more meaningful than preconditioners without a PDE term. In line 64 we mention a line of works that use KFAC for PINN, but only to approximate a preconditioner that does not include PDE terms. Hence, we regard these approaches as not satisfactory which motivated us to develop a KFAC method for preconditioners including PDE terms. We have slightly changed our wording in line 64 to improve clarity.\n\n> Line 131: Could the authors please elaborate on why the interior Gramian cannot be approximated with the existing K-FAC? Is this due to computational and/or autodiff limitations?\n\nK-FAC has only been proposed when the loss and therefore the pre-conditioner only incorporates function evaluations. This excludes PDE terms in the pre-conditioner, which appear naturally when incorporating second-order information in PINNs. Hence, the reason existing KFAC approaches can't be used is not a computational limitation but quite simply the lack of a KFAC for problems including PDE terms. \n\n> Is using a general loss in equation (11) necessary?\n\nWe decided to work with a general loss as losses other than the least-square problem appear in practice. This is for example the case for variational approaches including the [*deep Ritz method*](https://link.springer.com/article/10.1007/s40304-018-0127-z) and [*consistent PINNs*](https://arxiv.org/pdf/2406.09217) which  use Lp-norms. Of course, there is a trade-off between generality and accessibility. We decided to provide the general version as we already provide a didactic specific example with the Poisson equation and as we found the general form to be not much more complex. If you strongly disagree, we can also offer a simplified version in the main body and defer the fully general case to the appendix.\n\n> Does the proposed approach generalize to inverse problems involving PINNs?\n\nWe have not yet considered inverse problems. In general, we see no conceptual problem, however, one needs to take into account two (or more) neural networks when dealing with inverse problems. We leave this for future work.\n\n> Are there operators besides the Laplacian for which the weight sharing can be reduced?\n\nA similar reduction can be obtained for linear scalar valued PDE operators, meaning whenever the PDE of consideration is a linear function of the partial derivatives of some order.\n___\nThanks once more for your comments that have helped use improving our manuscript and we hope our responses could address your comments to your satisfaction. We remain attentive to your feedback!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qTRApXNp2v", "reviewer": "Reviewer_aiiA", "summary": "This work develops a KFAC approximation for approximate second order optimization of physics-informed neural network (PINN) losses, in order to address optimization difficulties of PINNs.The first key idea is to use forward mode automatic differentiation for the different derivatives of the PDE and notice that it corresponds to a forward pass with additional inputs (corresponding to the different derivatives) and weight sharing. The second key idea is to then tie this to KFAC approximations for weight sharing architectures. \nBoth of these ideas are already known previous work. However, the combination of these two ideas in the context of PINN optimization in order to obtain a scalable approximate second order method is novel and promising. Experimental evaluation involves only training loss curves for 2 PDEs (heat equation and poisson equation) as a function of time or iterations. This has been done only for a single run (no error bars), hyper-parameters were tuned with different HPO methods.", "questions": "- I noticed that the EMA update rate hyper-parameter is often super large, even up to 0.988 (for heat equation, appendix A9), which effectively looks only at the most recent batch. Do you have an explanation for why this is? What happens if you use a much smaller mini-batch size?\n- Is it correct that for linear PDEs, the approximation boils down to simply averaging over the additional forward passes corresponding to the additional derivatives on top of the average over the data points? So in comparison to the standard KFAC, you only have on additional average for these derivatives? \n- For the non-linear PDEs, you had to make an additional approximation in C26 - C27 of the supplementary material. Does this correspond to linearizing the non-linear PDE? Is this analogous to a last-layer non-linearity, e.g. if we had some additional tanh activation before the regression outputs in order to bound the output values between -1 and 1? If I understood C24 correctly, the Psi term appears just due to chain rule, but I did not understand why it was or had to be grouped with the Hessian of the loss w.r.t. outputs. Could you expand on this? \n- It seems very odd that KFAC performs better if the random mini-batch are *not* sampled (new) at every step. Have you noticed any such thing for KFAC on standard problems (not PINN losses)? Are the matrices in KFAC inverted every iteration or also just every T-th iteration?", "limitations": "The authors addressed the limitations, but I think the answer in (7) Experiment Statistical Significance should be No, as there are no error bars.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is generally very well written: The introduction, contributions and related work section is kept short and on the point, providing a clear explanation of the problem setting and motivating the approach of this paper.\n- The paper is technically solid. The authors are very precise in terms of the derivations and the notation. I could not spot any mistakes. \n- Viewing the derivatives as forward propagation of additional inputs with weight sharing and tying this to recently developed KFAC approximations is a neat insight, even if both on their own are not novel.", "weaknesses": "- While the authors are very rigorous in their notation (which I listed as a strength), it often took me a very long time to parse everything due to the overwhelming information for simple expressions. For instance, even the simple expression of the chain rule that has just two terms can look quite complicated when writing the Jacobians using an operator applied to a vector and both having multiple indices that correspond to indices for layer, data point, derivative. Again, this makes the math precise and avoids confusion. But I think that some parts could be simplified to help the reader to focus their attention on what is actually new throughout all the respective parts of the paper, rather than having to slowly parse every detail and remembering what each index and variable corresponds to. I am wondering if some of this could be simplified without much loss of precision. Perhaps, one quick win would be to remove layer indices (superscript l): Every layer is approximated independently (factorizing). The indices would only be important if you had cross-layer interactions, but I do not see this. Secondly, I wish you could write the Jacobians directly with a single symbol rather than operator applied to a vector or matrix (but I guess this is more difficult, because you need to differentiate between weights and neurons, etc.). \n\n- My main criticism is that the choice of experiments is unfortunately lacking similar rigor as the rest of the paper. The loss curves do show that the KFAC approximation is often favorable, but only for a single run. It would be important to do multiple runs to average out the noise from different initializations and mini-batching, and also include error bars. Furthermore, the loss curves only show the training loss and it is thus not clear if KFAC leads to more overfitting compared to other optimisers. I would have also liked to see some visualizations for some low-dimensional problems that show failure-cases of SGD and how KFAC fixes this. Or some other informative visualizations other than loss curves. \n\n---------------------\n- Minor: You noted that \"Taylor-mode\" AD is synonym to forward-mode AD and refer to related work in 3.1. However, \"Taylor-mode\" was mentioned a few times before, e.g. in the last sentence of related work, it is mentioned without a reference or explanation that it is also known as forward mode. Just put this information earlier.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer aiiA,\n\nthanks a lot for the time and effort you put into reviewing our work.\n\nRegarding readability, we agree that layer indices can be omitted and will do so in the updated version. We will also think about other ways to make the notation lighter.\n\n**Weaknesses**\n\n> [...] the loss curves only show the training loss and it is thus not clear if KFAC leads to more overfitting compared to other optimisers.\n\nWe respectfully disagree. *We report relative L2 error rather than training loss, which avoids overfitting*. The relative L2 error is defined as $\\frac{\\lVert u_\\theta - u^\\star\\rVert_{L^2(\\Omega)}}{\\lVert u^\\star\\rVert_{L^2(\\Omega)}}$, where $u^\\star$ is the true solution of the PDE. In PINNs, the relative L2 error is usually considered as the relevant quantity and we estimate it on held-out points that are not used during training, similar to the test loss in supervised learning.\n\nWe apologize for not explicitly introducing the relative L2 error and have adjusted this in our revised version.\n\n> The loss curves do show that the KFAC approximation is often favorable, but only for a single run. It would be important to do multiple runs to average out the noise from different initializations and mini-batching, and also include error bars.\n\nIt is true that our plots show the performance of a single run. However, this run was obtained through extensive hyper-parameter tuning, and we found our results to be consistent when using a different search strategy (random or Bayesian). We think this already supports the robustness of our results, but understand your concern. \n\nSince our computational ressources are limited, would you be satisfied if we presented an additional experiment with error bars for one of the sub-experiments?\n\n> I would have also liked to see some visualizations for some low-dimensional problems that show failure-cases of SGD and how KFAC fixes this. Or some other informative visualizations other than loss curves.\n\nThis is a great idea. We already have the code to do such visualizations because it was helpful to debug all optimizers. We will add such visualizations for the low-dimensional PDEs in the appendix.\n\n**Explicit questions**\n\n> I noticed that the EMA update rate hyper-parameter is often super large, even up to 0.988 (for heat equation, appendix A9), which effectively looks only at the most recent batch. Do you have an explanation for why this is? What happens if you use a much smaller mini-batch size? \n\nHigh EMA factors correspond to *slowly* forgetting pre-conditioners from previous iterations and are common in KFAC. For example, [Martens and Grosse](https://arxiv.org/pdf/1503.05671) uses an EMA factor increasing to 0.95 (page 19), which is also the [default value in KFAC's JAX implementation](https://github.com/google-deepmind/kfac-jax/blob/755f647d85252423f7dca6a14cf101735b5c46d8/kfac_jax/_src/optimizer.py#L298-L299).\n\n> Is it correct that for linear PDEs, the approximation boils down to simply averaging over the additional forward passes corresponding to the additional derivatives on top of the average over the data points? So in comparison to the standard KFAC, you only have on additional average for these derivatives?\n\nYes, this is correct and very concise. We have added a clarifying sentence using your suggested phrasing. \n\n> For the non-linear PDEs, you had to make an additional approximation in C26 - C27 of the supplementary material. [...] Could you expand on this?\n\nYes, the GN-preconditioner corresponds to linearizing the PDE. The approximations (C26)-(C27) are analoguous to the linear case and do not require further approximations compared to the linear case. \n\nThe interpretation of $\\Psi$ as a nonlinear output layer is correct. As $\\Psi$ is not part of the convex loss, we have to linearize it for the GN-preconditioner. We group the Jacobian of $\\Psi$ with the Hessian of the sample loss as it does not depend on the trainable parameters. It can also be grouped with the Jacobians $J_{n, \\alpha}^{(l)}$, however, this results in the same approximation. We hope that this clarifies the question, but please do not hesitate to let us know whether we should elaborate on this any further. \n\n> It seems very odd that KFAC performs better if the random mini-batch are not sampled (new) at every step. Have you noticed any such thing for KFAC on standard problems (not PINN losses)? Are the matrices in KFAC inverted every iteration or also just every T-th iteration?\n\nWe are not aware of any findings in the literature that have observed something like this. To make our comparison fair, we allowed the number of iterations a mini-batch is used as a hyper-parameter on which we performed Bayesian optimization. We found that for all optimizers it was favorable to keep a mini-batch for significant number of iterations. Thus, we regard this finding as something specific to the PINN problem rather than to KFAC. Our methods update the Kronecker matrices and inverses at every step. \n___\nThanks once more for your constructive comments. We hope our responses addressed them to your satisfaction. We remain attentive to your feedback!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We want to thank all reviewers for their thorough evaluation of our submission and provide extensive answers to the individual points raised by the reviewers below.

Here, we want to elaborate on our method's per-iteration complexity, both theoretically and empirically, as this concern was raised by multiple reviewers. We will add the following discussion to the main text.

On top of evaluating the loss and computing the gradient, there are two parts of KFAC that require additional resources; the assembly and inversion of the Kronecker factors: 

1. **Inversion of Kronecker factors:** Inverting layer $l$'s Kronecker approximation of the Gramian requires $O(h_l^3+h_{l+1}^3)$ time and $O(h_l^2+h_{l+1}^2)$ storage, where $h_l$ is the number of neurons in the $l$-th layer.  Note that inverting the exact block for layer $l$ (as done by 'ENGD (layer-wise)') requires $O(h_l^3h_{l+1}^3)$ time and $O(h_l^2 h_{l+1}^2)$ memory. In general, the improvement from the Kronecker factorization depends on how close to square the weight matrices of a layer are, and therefore on the architecture. In practise, the Kronecker factorization usually significantly reduces memory and run time.
Further improvements can be achieved by using structured Kronecker factors, e.g. diagonal or block-diagonal matrices as proposed in [1], which are beyond the scope of our paper.

2. **Assembly of Kronecker factors:** To compute the Kronecker factors, we use the layer inputs from evaluating the loss (for free), and need to backpropagate through its graph. The loss is computed with the forward Laplacian framework, which is the state-of-the-art method to compute input Laplacians of neural networks. The additional backpropagation is similar to the one performed by SGD with the only difference that we differentiate w.r.t. layer outputs rather than parameters.

We want to stress that evaluating the Laplacian and parameter derivatives is essential for *all* gradient-based PINN optimizers, not only for KFAC. To provide concrete insights into the computation time and memory of this routine, we ran an additional experiment and compared the forward Laplacian we use with the alternative `functorch` implementation. 

**We find that the forward Laplacian is roughly 2x faster and uses roughly half the memory compared to PyTorch's automatic differentiation. Further, the computation time for both methods scales linearly in the input dimension in our experiments.** 

We provide a detailed description of our findings below and in the attached PDF and will add it to the updated version.

**Conclusion**

Thanks again for your helpful comments and questions. We are convinced that they have helped us to substantially strengthen the manuscript and hope we addressed them to your satisfaction. 

We remain attentive to your feedback! 

## Experimental details

Due to the limited length of the discussion period, we present a short experiment. We will add a more detailed comparison to the manuscript.

**Setup:** We consider two Laplacian implementations:

1. **Autodiff Laplacian:** Computes the Laplacian with PyTorch's automatic differentiation (`functorch`) by computing the batched Hessian trace (via `torch.func.hessian` and `torch.func.vmap`). This is the standard approach in many PINN implementations.

2. **Forward Laplacian:** Computes the Laplacian via the forward Laplacian framework. We used this approach for all PDEs and optimizers presented in the experiments.

We use the biggest network from our experiments (the $D_\Omega \to 768 \to 768 \to 512 \to 512 \to 1$ MLP with tanh-activations from Fig. 3 right), then measure run time and peak memory consumption of computing the net's Laplacian on a mini-batch of size $N=1024$ with varying values of $D_\Omega$. To reduce measurement noise, we repeat each run over five independent Python sessions and report the smallest value. We use the same GPU as all other experiments, i.e. an NVIDIA RTX 6000 with 24 GiB memory.

**Results:** The following tables compare run time and peak memory between the two approaches:

| $D_\Omega$ | Autodiff Laplacian [s] | Forward Laplacian [s] |
|-----------|--------------------|-------------------|
| 1 | 0.051 (1.6x) | 0.033 (1.0x) |
| 10 | 0.20 (2.0x) | 0.10 (1.0x) |
| 100 | 1.7 (2.0x) | 0.84 (1.0x) |

We observe that the forward Laplacian is roughly twice as fast as the `functorch` Laplacian.

| $D_\Omega$ | Autodiff Laplacian [GiB] | Forward Laplacian [GiB] |
|-----------|--------------------|-------------------|
| 1 | 0.21 (0.96x) | 0.22 (1.0x) | 
| 10 | 0.98 (1.6x) | 0.61 (1.0x) |
| 100 | 8.8 (1.9x) | 4.6 (1.0x) |

We observe that the forward Laplacian uses significantly less memory for large input dimensions, up to only one half when $D_\Omega = 100$.

We visualized both tables using more values for $D_\Omega$ and observed linear scaling in both memory and run time; see the attached PDF. 

**Other scalings:**
- Both the forward Laplacian and nested backpropagation scale linearly w.r.t. batch size.
- Scaling w.r.t. output dimension (as asked by TQnM) is less relevant as many high-dimensional PDE problems are scalar-valued. 
- Scaling w.r.t. number of trainable parameters is highly architecture-dependent. This is also true for simple backpropagation in standard deep learning applications and we believe this question needs to be answered on a case-by-case basis. 

## References

[1] Lin, W., Dangel, F., Eschenhagen, R., Neklyudov, K., Kristiadi, A., Turner, R. E., & Makhzani, A. (ICML 2024). Structured inverse-free natural gradient descent: memory-efficient & numerically-stable KFAC.

Meta-review: All three reviewers agree this is a nice contribution to the literature of optimization for PINNs. The noted strengths are

* novel extension of KFAC to PINNs which leverages nontrivial forward-mode autodiff techniques. 
* rigorous and clear derivations
* empirical results demonstrating advantages over SGD, Adam, and L-BFGS

There are a few weaknesses such as computational complexity and robustness of experiments that are well addressed in rebuttal. Therefore, I'm recommending acceptance.
Decision: Accept (poster)