Title: Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: KFAC

PINNs

Gauss-Newton

PDEs

Taylor mode automatic differentiation

Forward Laplacian

Second-order optimization

Higher-order derivatives
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

Abstract

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

1 Introduction

Neural network-based approaches to numerically solve partial differential equations (PDEs) are growing at an unprecedented speed. The idea to train network parameters to minimize the residual of a PDE traces back to at least Dissanayake & Phan-Thien [15], Lagaris et al. [28], but was only recently popularized under the name deep Galerkin method (DGM) and Physics-informed neural networks (PINNs) through the works of Sirignano & Spiliopoulos [52], Raissi et al. [50]. PINNs are arguably one of the most popular network-based approaches to the numerical solution of PDEs as they are easy to implement, seamlessly incorporate measurement data, and promise to work well in high dimensions. Despite their immense popularity, PINNs are notoriously difficult to optimize [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

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 high-dimensional parameter space at cubic computational iteration cost, which prohibits scaling such approaches. To address this, we build on the idea of Kronecker-factored approximate curvature (KFAC) and apply it to Gauss-Newton matrices of PINN losses which greatly reduces the computational cost:

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

We use this weight sharing view to propose KFAC for Gauss-Newton matrices of objectives with differential operators, like PINN losses (§3.3 and 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 a 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 secondorder methods on small problems, scale more favorably to higher-dimensional neural networks and PDEs, and consistently outperform first-order methods and LBFGS (§4).

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

We build on the literature on Kronecker-factored approximate curvature (KFAC), which was initially introduced in Heskes [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 to 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 fullyconnected 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 first-index-varies-fastest convention, i.e. 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-valued objects back to the vector case by flattening a function's input and output before differentiation. The Jacobian of a vector-to-vector function a ↦→ b ( a ) has entries [J a b ] i,j = ∂b i / ∂a j . For a matrix-tomatrix function A ↦→ B ( A ) , the Jacobian is J A B = J vec A vec B . A useful property of vec is vec( AXB ) = ( B ⊤ ⊗ A ) vec X for matrices A , X , B which implies J X ( AXB ) = B ⊤ ⊗ A .

Sequential neural nets Consider a sequential neural network u θ = f θ ( L ) ◦ f θ ( L -1) ◦ . . . ◦ f θ (1) of depth L ∈ N . It consists of layers f θ ( l ) : R h ( l -1) → R h ( l ) , z ( l -1) ↦→ z ( l ) = f θ ( l ) ( z ( l -1) ) with

trainable parameters θ ( l ) ∈ R p ( l ) that transform an input z (0) := x ∈ R d := h (0) into a prediction u θ ( x ) = z ( L ) ∈ R h ( L ) via intermediate representations z ( l ) ∈ R h ( l ) . In the context of PINNs, we use networks with scalar outputs ( h ( L ) = 1 ) and denote the concatenation of all parameters by θ = ( θ (1) ⊤ , . . . , θ ( L ) ⊤ ) ⊤ ∈ R D . A common choice is to alternate fully-connected and activation layers. Linear layers map z ( l -1) ↦→ z ( l ) = W ( l ) z ( l -1) using a weight matrix W ( l ) = vec -1 θ ( l ) ∈ R h ( l ) × h ( l -1) (bias terms can be added as an additional column and by appending a 1 to the input). Activation layers map z ( l -1) ↦→ z ( l ) = σ ( z ( l -1) ) element-wise for a (typically smooth) σ : R → R .

2.1 Energy Natural Gradients for Physics-Informed Neural Networks

Let us consider a domain Ω ⊆ R d and the partial differential equation with right-hand side f , boundary data g and a differential operator L , e.g. the negative Laplacian -L u = ∆ x u = ∑ d i =1 ∂ 2 x i u . We parametrize u with a neural net and train its parameters θ to minimize the loss

with points { x n ∈ Ω } N Ω n =1 from the domain's interior, and points { x b n ∈ ∂ Ω } N ∂ Ω n =1 on its boundary. 2

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

Natural gradient methods perform parameter updates via a preconditioned gradient descent scheme θ ← θ -α G ( θ ) + ∇ L ( θ ) , where G ( θ ) + denotes the pseudo-inverse of a suitable Gramian matrix G ( θ ) ∈ R D × D and α is a step size. ENGD for the PINN loss (1) uses the Gramian

(2) is the Gauss-Newton matrix of the residual r ( θ ) = ( r Ω ( θ ) ⊤ / √ N Ω , r ∂ Ω ( θ ) ⊤ / √ N ∂ Ω ) ⊤ ∈ R N Ω + N ∂ Ω with interior and boundary residuals r Ω ,n ( θ ) = L u θ ( x n ) -f ( x n ) and r ∂ Ω ,n ( θ ) = u θ ( x b n ) -g ( x b n ) .

2.2 Kronecker-factored 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 / 2 N ∑ N n =1 ∥ u θ ( x n ) -y n ∥ 2 2 with targets y n ∈ R is

where u n = u θ ( x n ) , and it coincides with the classical Gauss-Newton matrix [37]. The established KFAC approximates (3). While the boundary Gramian G ∂ Ω ( θ ) has the same structure as F ( θ ) , the interior Gramian G Ω ( θ ) does not as it involves derivative rather than function evaluations of the net.

KFAC tackles the Fisher's per-layer block diagonal, F ( θ ) ≈ diag( F (1) ( θ ) , . . . , F ( L ) ( θ )) with F ( l ) ( θ ) = 1 / N ∑ N n =1 (J θ ( l ) u n ) ⊤ J θ ( l ) u n ∈ R p ( l ) × p ( l ) . For a fully-connected layer's block, let's examine the term J θ ( l ) u θ ( x ) from Equation (3) for a fixed data point. The layer parameters θ ( l ) = vec W ( l ) enter the computation via z ( l ) = W ( l ) z ( l -1) and we have J W ( l ) z ( l ) = z ( l -1) ⊤ ⊗ I [e.g. 10]. Further, the chain rule gives the decomposition J W ( l ) u = (J z ( l ) u )J W ( l ) z ( l ) = z ( l -1) ⊤ ⊗ J z ( l ) u . Inserting into F ( l ) ( θ ) , summing over data points, and using the expectation approximation ∑ n A n ⊗ B n ≈ N -1 ( ∑ n A n ) ⊗ ( ∑ n B n ) from Martens & Grosse [38], we obtain the KFAC approximation for linear layers in supervised square loss regression with a network's output,

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

3 Kronecker-Factored Approximate Curvature for PINNs

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

3.1 Higher-order Forward Mode Automatic Differentiation as Weight Sharing

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

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

For a linear layer f θ ( l ) ( z ( l -1) ) = W ( l ) z ( l -1) , applying the chain rule yields the propagation rule

The propagation rule through a nonlinear element-wise activation layer z ( l -1) ↦→ σ ( z ( l -1) ) is

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

This reduces computational cost, but is restricted to PDEs that involve second-order derivatives only via the Laplacian, or a partial Laplacian over a sub-set of input coordinates (e.g. heat equation, §4). For a more general second-order linear PDE operator L = ∑ d i,j =1 c i,j ∂ 2 x i ,x j , the forward pass for a linear layer is L z ( l ) = W ( l ) L z ( l -1) ∈ R h ( l ) , generalizing (7a), and similarly for Equation (7b)

see §C.3 for details. This is different from [30], which transforms the input space such that the coefficients are diagonal with entries { 0 , ± 1 } , reducing the computation to two forward Laplacians.

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

3.2 KFAC for Gauss-Newton Matrices with the Laplace Operator

Let's consider the Poisson equation's interior Gramian block for a linear layer (suppressing Ω in N Ω )

Because we made the Laplacian computation explicit through Taylor-mode autodiff (§3.1, specifically Equation (7a)), we can stack all output vectors that share the layer's weight into a matrix Z ( l ) n ∈ R h ( l ) × S with S = d + 2 and columns Z ( l ) n, 1 = z ( l ) n , Z ( l ) n, 2 = ∂ x 1 z ( l ) n , . . . , Z ( l ) n, 1+ d = ∂ x d z ( l ) n , and Z ( l ) n, 2+ d = ∆ x z ( l ) n (likewise Z ( l -1) n ∈ R h ( l -1) × S for the layer inputs), then apply the chain rule which has a structure similar to the Jacobian in §2.2, but with an additional sum over the S shared vectors. With that, we can now express the exact interior Gramian for a layer as

Next, we want to approximate Equation (8) with a Kronecker product. To avoid introducing a new convention, we rely on the KFAC approximation for linear layers with weight sharing developed by Eschenhagen et al. [17]-specifically, the approximation called KFAC-expand . This drops all terms with s = s ′ , then applies the expectation approximation from §2.2 over the batch and shared axes:

̸

KFAC for the Gauss-Newton matrix of a Laplace operator

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

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

where D m x u collects all partial derivatives of order m . For m ∈ { 0 , . . . , k } there are S m = ( d + m -1 d -1 ) independent partial derivatives and the total number of independent derivatives is S := ∑ k m =0 S m = ( d + k k ) . Ψ is a smooth mapping from all partial derivatives to R M , Ψ: R S → R M . To construct a PINN loss for Equation (10), we feed the residual r Ω ,n ( θ ) := Ψ( u θ ( x n ) , D x u θ ( x n ) , . . . , D k x u θ ( x n )) ∈ R M where D m x u θ ( x n ) ∈ R d × S m into a smooth convex criterion function ℓ : R M → R ,

The generalized Gauss-Newton (GGN) matrix [51] is the Hessian of L Ω ( θ ) when the residual is linearized w.r.t. θ before differentiation. It is positive semi-definite and has the form

with Λ ( r ) := ∇ 2 r ℓ ( r ) ∈ R M × M ≻ 0 the criterion's Hessian, e.g. ℓ ( r ) = 1 / 2 ∥ r ∥ 2 2 and Λ ( r ) = I M . Generalizing the second-order Taylor-mode from §3.1 to higher orders for the linear layer, we find

for any m . Hence, we can derive a forward propagation for the required derivatives where a linear layer processes at most S vectors 3 , i.e. the linear layer's weight is shared over the matrices D 0 x z ( l -1) := z ( l -1) , D 1 x z ( l -1) , . . . , D k x z ( l -1) . Stacking them into a matrix Z ( l -1) n = ( z ( l -1) , D 1 x z ( l -1) , . . . , D k x z ( l -1) ) ∈ R h ( l -1) × S (and Z ( l ) n for the outputs), the chain rule yields

where Z ( l -1) n,s ∈ R h ( l -1) denotes the s -th column of Z ( l -1) n . Following the same steps as in §3.2, we apply the KFAC-expand approximation from [17] to obtain the generalization of Equation (9):

KFAC for the GGN matrix of a general PDE operator

To bring this expression even closer to Equation (9), we can re-write the second Kronecker factor using an outer product decomposition Λ ( r Ω ,n ) = ∑ M m =1 l n,m l n,m with l n,m ∈ R M , then introduce g ( l ) n,s,m := (J Z ( l ) n,s r Ω ,n ) ⊤ l n,m ∈ R h ( l ) and write the second term as 1 / N ∑ N,S,M n,s,m =1 g ( l ) n,s,m g ( l ) ⊤ n,s,m , similar to the Kronecker-factored low-rank (KFLR) approach of Botev et al. [5].

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

where ℓ : R K → R is a convex function. We can perceive this as a special case of the setting above with Ψ = 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. §B shows pseudo-code for our method on the Poisson equation.

At iteration t , we approximate the per-layer interior and boundary Gramians using our derived Kronecker approximation (Equations (9) and (14)), G ( l ) Ω ,t ≈ A ( l ) Ω ,t ⊗ B ( l ) Ω ,t and G ( l ) ∂ Ω ,t ≈ A ( l ) Ω ,t ⊗ B ( l ) Ω ,t .

Exponential moving average and damping For preconditioning, we accumulate the Kronecker factors A ( l ) · ,t , B ( l ) · ,t over time using an exponential moving average ˆ A ( l ) · ,t = β ˆ A ( l ) · ,t -1 +(1 -β ) A ( l ) · ,t of factor β ∈ [0 , 1) (identically for ˆ B ( l ) · ,t ), similar to the original KFAC. Moreover, we apply the same constant damping of strength λ > 0 to all Kronecker factors, ˜ A ( l ) · ,t = ˆ A ( l ) · ,t + λ I and ˜ B ( l ) · ,t = ˆ B ( l ) · ,t + λ I such that the curvature approximation used for preconditioning at step t is

Gradient preconditioning Given layer l 's mini-batch gradient g ( l ) t = ∂L ( θ t ) / ∂ θ ( l ) t ∈ R p ( l ) , we obtain an update direction ∆ ( l ) t = -( ˜ A ( l ) Ω ,t ⊗ ˜ B ( l ) Ω ,t + ˜ A ( l ) ∂ Ω ,t ⊗ ˜ B ( l ) ∂ Ω ,t ) -1 g ( l ) t ∈ R p ( l ) using the trick of [38, Appendix I] to invert the Kronecker sum via eigen-decomposing all Kronecker factors.

Learning rate and momentum From the preconditioned gradient ∆ t ∈ R D , we consider two different updates θ t +1 = θ t + δ t we call KFAC and KFAC* . KFAC uses momentum over previous updates, ˆ δ t = µ δ t -1 + ∆ t , and µ is chosen by the practitioner. Like ENGD, it uses a logarithmic grid line search, selecting δ t = α ⋆ ˆ δ t with α ⋆ = arg min α L ( θ t + α ˆ δ t ) where α ∈ { 2 -30 , . . . , 2 0 } . KFAC* uses the automatic learning rate and momentum heuristic of the original KFAC optimizer. It parametrizes the iteration's update as δ t +1 ( α, µ ) = α ∆ t + µ δ t , then obtains the optimal parameters by minimizing the quadratic model m ( δ t +1 ) = L ( θ t ) + δ ⊤ t +1 g t + 1 / 2 δ ⊤ t +1 ( G ( θ t ) + λ I ) δ t +1 with the exact damped Gramian. The optimal learning rate and momentum arg min α,µ m ( δ t +1 ) are

(see [38, Section 7] for details). The computational cost is dominated by the two Gramian-vector products with ∆ t and δ t . By using the Gramian's outer product structure [12, 45], we perform them with autodiff [48, 51] using one Jacobian-vector product each, as recommended in [38].

Computational complexity 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, whereas inverting the exact block for layer would require O ( h ( l ) 3 h ( 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. (block-)diagonal 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.

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

2

10

1

L

2

error

D

= 257

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

D

= 9

,

873

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

L

2

error

D

= 449

10

0

10

1

10

2

10

3

Time [s]

10

-

4

10

-

2

10

0

D

= 10

,

065

10

0

10

1

10

2

10

3

Time [s]

10

-

4

10

-

2

10

0

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

4 Experiments

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

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

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

For the shallow net (Figure 1, left), we can reproduce the results of [41], where exact ENGD achieves high accuracy. In terms of computation time, our KFACs are competitive with full-ENGD for a long

10

0

10

1

10

2

10

3

Time [s]

10

-

3

10

-

1

10

1

L

2

error

5d Poisson,

D

= 116

,

865

10

0

10

1

10

2

10

3

Time [s]

10

-

3

10

-

1

10d Poisson,

D

= 118

,

145

10

0

10

2

10

4

Time [s]

10

-

2

10

-

1

10

0

100d Poisson,

D

= 1

,

325

,

057

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

phase, outperforming the first-order and quasi-Newton baselines. In contrast to ENGD, which runs out of memory for networks with more than 10 000 parameters, KFAC scales to larger networks (Figure 1, center and right) and is competitive with other second-order optimizers like Hessian-free, which uses more sophisticated heuristics. We make similar observations on a small (1+1)d heat equation with the same models, see §A.7 and 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 a four-dimensional heat equation

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 in §A.8), see Figure 2. To prevent over-fitting, we use mini-batches and sample a new batch each iteration. We noticed that KFAC improves significantly when batches are sampled less frequently and hypothesize that it might need more iterations to make similar progress than one iteration of Hessian-free or ENGD on a batch. Consequently we sample a new batch only every 100 iterations for KFAC. To ensure that this does not lead to an unfair advantage for KFAC, we conduct an additional experiment for the MLP with D = 116864 where we tune batch sizes, batch sampling frequencies, and all hyper-parameters with generous search spaces using Bayesian search (§A.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 on the same architectures, see §A.3 and fig. A7.

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

We use the same architectures as before, but with larger intermediate widths and parameters up to a million (Figure 3). Due to lacking references for training such high-dimensional problems, we

select all hyper-parameters via Bayesian search, including batch sizes and batch sampling frequencies (details in §A.5). We see a similar picture as before with KFAC consistently outperforming firstorder methods and LBFGS, offering competitive performance with Hessian-free. To account for the possibility that the Bayesian search did not properly identify good hyper-parameters, we conduct a random/grid search experiment for the 10d Poisson equation (Figure 3, middle), using similar batch sizes and same batch sampling frequencies as for the (4 + 1) d heat equation (details in §A.4). In this experiment, KFAC also achieved similar performance than Hessian-free and outperformed SGD, Adam, and LBFGS (Figure A8).

(9+1)d Fokker-Planck equation To show the applicability to nonlinear PDEs, we consider a Fokker-Planck equation in logarithmic space. PINN formulations of the Fokker-Planck equation have been considered in [23, 54]. Concretely, we are solving a nine-dimensional equation of the form

with d = 9 , t ∈ [0 , 1] and x ∈ R 9 , where in practice R 9 is replaced by [ -5 , 5] 9 . The solution is q ∗ = log( p ∗ ) and p ∗ is given by p ∗ ( t, x ) ∼ N (0 , exp( -t ) I +(1 -exp( -t ))2 I ) . Wemodel the solution with a medium sized tanh-activated MLP with D = 118145 parameters, batch sizes are N Ω = 3000 , N ∂ Ω = 1000 , and we assign each run a computation time budget of 6 000 s. As in previous experiments, the batches are resampled every iteration for all optimizers except for KFAC and KFAC*, which use the same batch for ten steps (details in §A.11). Figure 4 reports the L 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.

5 Discussion and Conclusion

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

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

10

0

10

1

10

2

10

3

Time [s]

10

-

1

10

0

10

1

L

2

error

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

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 Forschungsgemeinschaft (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

Amari, S.-I. Natural gradient works efficiently in learning. Neural computation , 10(2):251-276, 1998.

Benzing, F. Gradient Descent on Neurons and its Link to Approximate Second-order Optimization. In International Conference on Machine Learning (ICML) , 2022.

Bettencourt, J., Johnson, M. J., and Duvenaud, D. Taylor-mode automatic differentiation for higher-order derivatives in JAX. In Advances in Neural Information Processing Systems (NeurIPS); Workhop on Program Transformations for ML , 2019.

Bonfanti, A., Bruno, G., and Cipriani, C. The Challenges of the Nonlinear Regime for PhysicsInformed Neural Networks. arXiv preprint arXiv:2402.03864 , 2024.

Botev, A., Ritter, H., and Barber, D. Practical Gauss-Newton optimisation for deep learning. In International Conference on Machine Learning (ICML) , 2017.

Bottou, L., Curtis, F. E., and Nocedal, J. Optimization methods for large-scale machine learning. SIAM Review (SIREV) , 60, 2016.

Chen, Z., McCarran, J., Vizcaino, E., Soljaˇ ci´ c, M., and Luo, D. TENG: Time-Evolving Natural Gradient for Solving PDEs with Deep Neural Net. arXiv preprint arXiv:2404.10771 , 2024.

Clarke, R. M., Su, B., and Hernández-Lobato, J. M. Adam through a second-order lens. 2023.

Cuomo, S., Di Cola, V. S., Giampaolo, F., Rozza, G., Raissi, M., and Piccialli, F. Scientific machine learning through physics-informed neural networks: Where we are and what's next. Journal of Scientific Computing , 92(3):88, 2022.

Dangel, F., Harmeling, S., and Hennig, P. Modular Block-diagonal Curvature Approximations for Feedforward Architectures. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2020.

Dangel, F., Kunstner, F., and Hennig, P. BackPACK: Packing more into Backprop. In International Conference on Learning Representations (ICLR) , 2020.

Dangel, F., Tatzel, L., and Hennig, P. ViViT: Curvature access through the generalized gaussnewton's low-rank structure. Transactions on Machine Learning Research (TMLR) , 2022.

Daw, A., Bu, J., Wang, S., Perdikaris, P., and Karpatne, A. Rethinking the importance of sampling in physics-informed neural networks. arXiv preprint arXiv:2207.02338 , 2022.

De Ryck, T., Bonnet, F., Mishra, S., and de Bézenac, E. An operator preconditioning perspective on training in physics-informed machine learning. arXiv preprint arXiv:2310.05801 , 2023.

Dissanayake, M. and Phan-Thien, N. Neural-network-based approximations for solving partial differential equations. communications in Numerical Methods in Engineering , 10(3):195-201, 1994.

E, W. and Yu, B. The deep ritz method: a deep learning-based numerical algorithm for solving variational problems. Communications in Mathematics and Statistics , 6(1):1-12, 2018.

Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., and Hennig, P. Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures. In Advances in Neural Information Processing Systems (NeurIPS) , 2023.

Griewank, A. and Walther, A. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

Griewank, A., Juedes, D., and Utke, J. Algorithm 755: ADOL-C: A package for the automatic differentiation of algorithms written in C/C++. ACM Transactions on Mathematical Software (TOMS) , 22(2):131-167, 1996.

Grosse, R. and Martens, J. A Kronecker-Factored Approximate Fisher Matrix for Convolution Layers. In International Conference on Machine Learning (ICML) , 2016.

Grosse, R., Bae, J., Anil, C., Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E., et al. Studying large language model generalization with influence functions. arXiv preprint arXiv:2308.03296 , 2023.

Heskes, T. On 'natural' learning and pruning in multilayered perceptrons. Neural Computation , 12(4):881-901, 2000.

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.

Jnini, A., Vella, F., and Zeinhofer, M. Gauss-Newton Natural Gradient Descent for PhysicsInformed Computational Fluid Dynamics. arXiv preprint arXiv:2402.10680 , 2024.

Johnson, M. J., Bettencourt, J., Maclaurin, D., and Duvenaud, D. Taylor-made higher-order automatic differentiation. 2021. URL https://github.com/google/jax/files/6717197/ jet.pdf . Accessed January 03, 2024.

Krishnapriyan, A., Gholami, A., Zhe, S., Kirby, R., and Mahoney, M. W. Characterizing possible failure modes in physics-informed neural networks. Advances in Neural Information Processing Systems , 34:26548-26560, 2021.

Kunstner, F., Hennig, P., and Balles, L. Limitations of the empirical Fisher approximation for natural gradient descent. Advances in neural information processing systems , 32, 2019.

Lagaris, I. E., Likas, A., and Fotiadis, D. I. Artificial neural networks for solving ordinary and partial differential equations. IEEE transactions on neural networks , 9(5):987-1000, 1998.

Li, R., Ye, H., Jiang, D., Wen, X., Wang, C., Li, Z., Li, X., He, D., Chen, J., Ren, W., et al. Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlo. 2023.

Li, R., Wang, C., Ye, H., He, D., and Wang, L. DOF: Accelerating high-order differential operators with forward propagation. In International Conference on Learning Representations (ICLR), Workshop on AI4DifferentialEquations In Science , 2024.

Lin, W., Duruisseaux, V., Leok, M., Nielsen, F., Khan, M. E., and Schmidt, M. Simplifying Momentum-based Riemannian Submanifold Optimization. 2023.

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.

Liu, S., Su, C., Yao, J., Hao, Z., Su, H., Wu, Y ., and Zhu, J. Preconditioning for physics-informed neural networks. arXiv preprint arXiv:2402.00531 , 2024.

Lu, L., Meng, X., Mao, Z., and Karniadakis, G. E. DeepXDE: A deep learning library for solving differential equations. SIAM Review , 63(1):208-228, 2021.

Markidis, S. The old and the new: Can physics-informed deep-learning replace traditional linear solvers? Frontiers in big Data , 4:669097, 2021.

Martens, J. Deep learning via Hessian-free optimization. In International Conference on Machine Learning (ICML) , 2010.

Martens, J. New insights and perspectives on the natural gradient method, 2020.

Martens, J. and Grosse, R. Optimizing Neural Networks with Kronecker-factored Approximate Curvature. In International Conference on Machine Learning (ICML) , 2015.

Martens, J., Ba, J., and Johnson, M. Kronecker-factored Curvature Approximations for Recurrent Neural Networks. In International Conference on Learning Representations , 2018. URL https://openreview.net/forum?id=HyMTkQZAb .

Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., et al. Mixed precision training. 2017.

Müller, J. and Zeinhofer, M. Achieving high accuracy with PINNs via energy natural gradient descent. In International Conference on Machine Learning , pp. 25471-25485. PMLR, 2023.

Müller, J. and Zeinhofer, M. Optimization in SciML-A Function Space Perspective. arXiv preprint arXiv:2402.07318 , 2024.

Nabian, M. A., Gladstone, R. J., and Meidani, H. Efficient training of physics-informed neural networks via importance sampling. Computer-Aided Civil and Infrastructure Engineering , 36 (8):962-977, 2021.

Osawa, K., Li, S., and Hoefler, T. Pipefisher: Efficient training of large language models using pipelining and Fisher information matrices. Proceedings of Machine Learning and Systems , 5, 2023.

Papyan, V. Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet Hessians. In International Conference on Machine Learning (ICML) , 2019.

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems (NeurIPS) . 2019.

Pauloski, J. G., Huang, Q., Huang, L., Venkataraman, S., Chard, K., Foster, I., and Zhang, Z. Kaisa: an adaptive second-order optimizer framework for deep neural networks. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis , pp. 1-14, 2021.

Pearlmutter, B. A. Fast Exact Multiplication by the Hessian. Neural Computation , 1994.

Petersen, F., Sutter, T., Borgelt, C., Huh, D., Kuehne, H., Sun, Y., and Deussen, O. ISAAC Newton: Input-based Approximate Curvature for Newton's Method. In International Conference on Learning Representations (ICLR) , 2023.

Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics , 378:686-707, 2019.

Schraudolph, N. N. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation , 2002.

Sirignano, J. and Spiliopoulos, K. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics , 375:1339-1364, 2018.

Skorski, M. Chain rules for hessian and higher derivatives made easy by tensor calculus. arXiv preprint arXiv:1911.13292 , 2019.

Sun, J., Berner, J., Richter, L., Zeinhofer, M., Müller, J., Azizzadenesheli, K., and Anandkumar, A. Dynamical measure transport and neural pde solvers for sampling. arXiv preprint arXiv:2407.07873 , 2024.

Tatzel, L., Hennig, P., and Schneider, F. Late-Phase Second-Order Training. In Advances in Neural Information Processing Systems (NeurIPS), Workshop Has it Trained Yet? , 2022.

van der Meer, R., Oosterlee, C. W., and Borovykh, A. Optimally weighted loss functions for solving PDEs with neural networks. Journal of Computational and Applied Mathematics , 405: 113887, 2022.

Wang, S., Teng, Y., and Perdikaris, P. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal on Scientific Computing , 43(5):A3055-A3081, 2021.

Wang, S., Sankaran, S., and Perdikaris, P. Respecting causality is all you need for training physics-informed neural networks. arXiv preprint arXiv:2203.07404 , 2022.

Wang, S., Yu, X., and Perdikaris, P. When and why PINNs fail to train: A neural tangent kernel perspective. Journal of Computational Physics , 449:110768, 2022.

Weights and Biases. Experiment Tracking with Weights and Biases, 2020. URL https: //www.wandb.ai/ . Software available from wandb.ai.

Wu, C., Zhu, M., Tan, Q., Kartha, Y., and Lu, L. A comprehensive study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering , 403:115671, 2023.

Zampini, S., Zerbinati, U., Turkiyyah, G., and Keyes, D. PETScML: Second-order solvers for training regression problems in Scientific Machine Learning. arXiv preprint arXiv:2403.12188 , 2024.

Zapf, B., Haubner, J., Kuchta, M., Ringstad, G., Eide, P. K., and Mardal, K.-A. Investigating molecular transport in the human brain from MRI with physics-informed neural networks. Scientific Reports , 12(1):1-12, 2022.

Zeng, Q., Kothari, Y., Bryngelson, S. H., and Schäfer, F. Competitive physics informed networks. arXiv preprint arXiv:2204.11144 , 2022.

Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks (Supplementary Material)

A Experimental Details and Additional Results

A.1 Hyper-Parameter Tuning Protocol

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

SGD: learning rate, momentum

Adam: learning rate

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

LBFGS: learning rate, history size

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

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

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

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

A.2 2d Poisson Equation

Setup We consider a two-dimensional Poisson equation -∆ u ( x, y ) = 2 π 2 sin( πx ) sin( πy ) 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 , 1] 2 . We choose a single set of training points with N Ω = 900 , N ∂ Ω = 120 . The L 2 error is evaluated on a separate set of 9 000 data points using the known solution u ⋆ ( x, y ) = sin( πx ) sin( πy ) . Each run is limited to a compute time of 1 000 s. We compare three MLP architectures of increasing size, each of whose linear layers are Tanh-activated except for the final one: a shallow 2 → 64 → 1 MLP with D = 257 trainable parameters, a five layer 2 → 64 → 64 → 48 → 48 → 1 MLP with D = 9873 trainable parameters, and a five layer 2 → 256 → 256 → 128 → 128 → 1 MLP with D = 116097 trainable parameters. For the biggest architecture, full and per-layer ENGD lead to out-of-memory errors and are thus not tested in the experiments. Figure A5 visualizes the results, and Figure A6 illustrates the learned solutions over training for all optimizers on the shallow MLP.

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

2 → 64 → 1 MLP with D = 257

Adam: learning rate: 1 . 692 339 · 10 -3

SGD: learning rate: 1 . 805 015 · 10 -2 , momentum: 9 . 9 · 10 -1

-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: 1 · 10 -10 , exponential moving average: 3 · 10 -1 , initialize Gramian to identity: yes

-ENGD (layer-wise): damping: 0 , exponential moving average: 9 · 10 -1 , initialize Gramian to identity: yes

KFAC: damping: 1 . 544 099 · 10 -12 , momentum: 5 . 117 575 · 10 -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 → 1 MLP with D = 9873

Adam: learning rate: 2 . 052 448 · 10 -4

SGD: learning rate: 3 . 758 303 · 10 -3 , momentum: 9 · 10 -1

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

2

10

1

L

2

error

D

= 257

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

D

= 9

,

873

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

10

10

-

6

10

-

2

10

2

Loss

D

= 257

10

0

10

1

10

2

10

3

Time [s]

10

-

8

10

-

4

10

0

D

= 9

,

873

10

0

10

1

10

2

10

3

Time [s]

10

-

7

10

-

3

10

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

(b)

10

1

10

3

10

5

Iteration

10

-

5

10

-

2

10

1

L

2

error

D

= 257

10

1

10

3

10

5

Iteration

10

-

5

10

-

3

10

-

1

D

= 9

,

873

10

1

10

3

10

5

Iteration

10

-

5

10

-

3

10

-

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

1

10

3

10

5

Iteration

10

-

10

10

-

6

10

-

2

10

2

Loss

D

= 257

10

1

10

3

10

5

Iteration

10

-

8

10

-

4

10

0

D

= 9

,

873

10

1

10

3

10

5

Iteration

10

-

7

10

-

3

10

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

-Hessian-free: curvature matrix: GGN, initial damping: 1 · 10 -1 , constant damping: no, maximum CG iterations: 350

LBFGS: learning rate: 1 · 10 -1 , history size: 200

-ENGD (full): damping: 1 · 10 -10 , exponential moving average: 6 · 10 -1 , initialize Gramian to identity: no

-ENGD(layer-wise): damping: 1 · 10 -6 , exponential moving average: 3 · 10 -1 , initialize Gramian to identity: no

KFAC: damping: 2 . 640 390 · 10 -11 , momentum: 9 . 995 595 · 10 -2 , exponential moving average: 5 . 556 664 · 10 -1 , 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 → 1 MLP with D = 116097

Adam: learning rate: 6 . 406 108 · 10 -4

SGD: learning rate: 2 . 478 674 · 10 -3 , momentum: 9 · 10 -1

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

LBFGS: learning rate: 1 · 10 -1 , history size: 225

Optimizer

First step

0.1% trained

1%trained

10% trained

True solution

SGD

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

5

.

3

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

335

, Loss:

2

.

8

·

10

-

1

,

L

2

loss:

2

.

4

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

3299

, Loss:

1

.

7

·

10

-

4

,

L

2

loss:

3

.

6

·

10

-

3

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

32494

, Loss:

5

.

6

·

10

-

6

,

L

2

loss:

1

.

8

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

5

.

3

·

10

-

1

Adam

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

8

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

335

, Loss:

5

.

6

·

10

0

,

L

2

loss:

4

.

1

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

3299

, Loss:

1

.

4

·

10

-

2

,

L

2

loss:

4

.

1

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

32494

, Loss:

9

.

8

·

10

-

6

,

L

2

loss:

7

.

7

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

8

·

10

-

1

LBFGS

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

4

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

50

, Loss:

4

.

0

·

10

-

5

,

L

2

loss:

1

.

1

·

10

-

3

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

491

, Loss:

1

.

0

·

10

-

6

,

L

2

loss:

1

.

8

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

4831

, Loss:

1

.

0

·

10

-

6

,

L

2

loss:

1

.

8

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

4

·

10

-

1

Hessian-free

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

1

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

4

, Loss:

3

.

9

·

10

1

,

L

2

loss:

1

.

8

·

10

0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

35

, Loss:

1

.

0

·

10

-

4

,

L

2

loss:

3

.

1

·

10

-

3

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

335

, Loss:

2

.

1

·

10

-

8

,

L

2

loss:

1

.

8

·

10

-

5

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

1

·

10

-

1

ENGD (full)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

1

.

5

·

10

0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

6

, Loss:

1

.

1

·

10

0

,

L

2

loss:

3

.

4

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

55

, Loss:

7

.

8

·

10

-

1

,

L

2

loss:

1

.

3

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

594

, Loss:

5

.

1

·

10

-

3

,

L

2

loss:

1

.

9

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

1

.

5

·

10

0

ENGD

(layer-wise)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

7

·

10

0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

6

, Loss:

1

.

7

·

10

0

,

L

2

loss:

3

.

7

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

61

, Loss:

8

.

0

·

10

-

3

,

L

2

loss:

3

.

3

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

594

, Loss:

3

.

1

·

10

-

11

,

L

2

loss:

6

.

6

·

10

-

7

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

3

.

7

·

10

0

KFAC

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

1

.

7

·

10

0

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

20

, Loss:

3

.

4

·

10

-

1

,

L

2

loss:

3

.

3

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

208

, Loss:

4

.

9

·

10

-

6

,

L

2

loss:

1

.

5

·

10

-

3

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

2049

, Loss:

8

.

4

·

10

-

10

,

L

2

loss:

1

.

1

·

10

-

5

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

1

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

1

.

7

·

10

0

KFAC*

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

4

.

0

·

10

-

1

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

Normalized true solution

Step:

35

, Loss:

7

.

2

·

10

-

2

,

L

2

loss:

7

.

5

·

10

-

2

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

Normalized true solution

Step:

335

, Loss:

5

.

7

·

10

-

5

,

L

2

loss:

2

.

1

·

10

-

3

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

Normalized true solution

Step:

3629

, Loss:

1

.

8

·

10

-

9

,

L

2

loss:

4

.

4

·

10

-

6

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

2

Normalized learned solution

Normalized true solution

Step:

0

, Loss:

4

.

9

·

10

1

,

L

2

loss:

4

.

0

·

10

-

1

0

.

8

0

0

.

.

6

8

1

.

0

0

1

.

.

8

0

0

.

0

1

.

0

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

0

.

0

0

.

0

0

.

0

0

.

0

0

.

0

.

4

0

.

6

0

.

6

0

.

8

0

.

8

1

.

0

1

.

0

0

.

0

0

.

0

0

.

2

0

0

.

.

0

2

0

.

4

0

0

.

.

2

4

0

.

6

0

0

.

.

4

6

0

1

.

0

0

.

0

0

.

0

0

.

2

0

.

2

0

.

4

0

.

4

0

.

6

0

.

6

0

.

8

0

.

8

1

.

0

1

.

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

.

0

0

.

0

0

.

2

0

.

2

0

.

4

0

KFAC: damping: 1 . 710 269 · 10 -13 , momentum: 8 . 484 996 · 10 -1 , exponential moving average: 9 . 636 460 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC*: damping: 1 . 232 407 · 10 -13 , exponential moving average: 9 . 488 207 · 10 -1 , initialize Kronecker factors to identity: yes

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

2 → 64 → 1 MLP with D = 257

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -1 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 100 , 1 , 1 · 10 -2 , 1 · 10 -4 , 1 · 10 -6 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 75 , 100 , 125 , 150 , 175 , 200 , 225 , 250 } )

-ENGD(layer-wise): damping: U ( { 1 · 10 -4 , 1 · 10 -6 , 1 · 10 -8 , 1 · 10 -10 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

-ENGD (full): damping: U ( { 1 · 10 -6 , 1 · 10 -8 , 1 · 10 -10 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC*: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

2 → 64 → 64 → 48 → 48 → 1 MLP with D = 9873

Adam: learning rate: LU ([1 · 10 -4 ; 5 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

-ENGD (full): damping: U ( { 1 · 10 -8 , 1 · 10 -9 , 1 · 10 -10 , 1 · 10 -11 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

-ENGD (layer-wise): damping: U ( { 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 , 1 · 10 -5 , 1 · 10 -6 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC*: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } )

2 → 256 → 256 → 128 → 128 → 1 MLP with D = 116097

Adam: learning rate: LU ([1 · 10 -4 ; 5 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

KFAC*: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } )

A.3 5d Poisson Equation

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

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

2

10

1

L

2

error

D

= 449

10

0

10

1

10

2

10

3

Time [s]

10

-

3

10

-

1

D

= 10

,

065

10

0

10

1

10

2

10

3

Time [s]

10

-

3

10

-

1

10

1

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

8

10

-

3

10

2

Loss

D

= 449

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

2

10

1

D

= 10

,

065

10

0

10

1

10

2

10

3

Time [s]

10

-

4

10

0

10

4

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

(b)

10

1

10

3

10

5

Iteration

10

-

5

10

-

2

10

1

L

2

error

D

= 449

10

1

10

3

10

5

Iteration

10

-

3

10

-

1

D

= 10

,

065

10

1

10

3

Iteration

10

-

3

10

-

1

10

1

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

1

10

3

10

5

Iteration

10

-

8

10

-

3

10

2

Loss

D

= 449

10

1

10

3

10

5

Iteration

10

-

5

10

-

2

10

1

D

= 10

,

065

10

1

10

3

Iteration

10

-

4

10

0

10

4

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

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

5 → 64 → 1 MLP with D = 449

Adam: learning rate: 1 . 527 101 · 10 -3

SGD: learning rate: 4 . 829 757 · 10 -3 , momentum: 9 · 10 -1

-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 -6 , exponential moving average: 6 · 10 -1 , initialize Gramian to identity: yes

KFAC: damping: 3 . 030 734 · 10 -13 , momentum: 4 . 410 155 · 10 -1 , exponential moving average: 3 . 260 663 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC*: damping: 9 . 835 853 · 10 -13 , exponential moving average: 7 . 714 287 · 10 -1 , initialize Kronecker factors to identity: yes

5 → 64 → 64 → 48 → 48 → 1 MLP with D = 10065

Adam: learning rate: 6 . 999 994 · 10 -4

SGD: learning rate: 1 . 007 555 · 10 -3 , momentum: 9 · 10 -1

-Hessian-free: curvature matrix: GGN, initial damping: 1 · 10 -1 , constant damping: no, maximum CG iterations: 350

LBFGS: learning rate: 1 · 10 -1 , history size: 225

-ENGD (full): damping: 1 · 10 -8 , exponential moving average: 3 · 10 -1 , initialize Gramian to identity: no

-ENGD(layer-wise): damping: 1 · 10 -6 , exponential moving average: 3 · 10 -1 , initialize Gramian to identity: no

KFAC: damping: 1 . 183 063 · 10 -12 , momentum: 9 . 058 900 · 10 -1 , 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 to identity: yes

5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116865

Adam: learning rate: 5 . 416 376 · 10 -4

SGD: learning rate: 1 . 924 173 · 10 -3 , momentum: 9 · 10 -1

-Hessian-free: curvature matrix: GGN, initial damping: 5 , constant damping: no, maximum CG iterations: 350

LBFGS: learning rate: 2 · 10 -2 , history size: 225

KFAC*: damping: 2 . 183 605 · 10 -12 , exponential moving average: 9 . 563 992 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC: damping: 1 . 844 213 · 10 -11 , momentum: 7 . 528 559 · 10 -1 , exponential moving average: 9 . 307 849 · 10 -1 , initialize Kronecker factors to identity: yes

Search space details The runs shown in Figure A7 were determined to be the best via a search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search ( U denotes a uniform, and LU a log-uniform distribution):

5 → 64 → 1 MLP with D = 449

Adam: learning rate: LU ([1 · 10 -4 ; 5 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

-ENGD (full): damping: U ( { 1 · 10 -8 , 1 · 10 -9 , 1 · 10 -10 , 1 · 10 -11 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

-ENGD (layer-wise): damping: U ( { 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 , 1 · 10 -5 , 1 · 10 -6 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC*: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

5 → 64 → 64 → 48 → 48 → 1 MLP with D = 10065

Adam: learning rate: LU ([1 · 10 -4 ; 5 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

-ENGD (full): damping: U ( { 1 · 10 -8 , 1 · 10 -9 , 1 · 10 -10 , 1 · 10 -11 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

-ENGD (layer-wise): damping: U ( { 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 , 1 · 10 -5 , 1 · 10 -6 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC*: damping: LU ([1 · 10 -12 ; 1 · 10 -6 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116865

Adam: learning rate: LU ([1 · 10 -4 ; 5 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

KFAC*: damping: LU ([1 · 10 -12 ; 1 · 10 -6 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

A.4 10d Poisson Equation

Setup We consider a 10-dimensional Poisson equation -∆ u ( x ) = 0 on the 10-dimensional unit square x ∈ [0 , 1] 5 with zero right-hand side and harmonic mixed second order polynomial boundary conditions u ( x ) = ∑ d / 2 i =1 x 2 i -1 x 2 i for x ∈ ∂ [0 , 1] d . We sample training batches of size N Ω = 3000 , N ∂ Ω = 1000 and evaluate the L 2 error on a separate set of 30 000 data points using the known solution u ⋆ ( x ) = ∑ d / 2 i =1 x 2 i -1 x 2 i . All optimizers except for KFAC sample a new training batch each iteration. KFAC only re-samples every 100 iterations because we noticed significant improvement with multiple iterations on a fixed batch. Each run is limited to 6 000 s. We use a 10 → 256 → 256 → 128 → 128 → 1 MLP with D = 118145 MLP whose linear layers are Tanh-activated except for the final one. Figure A8 visualizes the results.

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

4

10

-

3

10

-

2

10

-

1

10

0

L

2

error

SGD

Adam

Hessian-free

LBFGS

KFAC

10

0

10

1

10

2

10

3

Time [s]

10

-

8

10

-

6

10

-

4

10

-

2

10

0

Loss

(b)

10

0

10

1

10

2

10

3

10

4

10

5

Iteration

10

-

4

10

-

3

10

-

2

10

-

1

10

0

L

2

error

10

0

10

1

10

2

10

3

10

4

10

5

Iteration

10

-

8

10

-

6

10

-

4

10

-

2

10

0

Loss

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

SGD: learning rate: 6 . 550 109 · 10 -3 , momentum: 9 · 10 -1

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

Adam: learning rate: 1 . 359 480 · 10 -4

LBFGS: learning rate: 2 · 10 -1 , history size: 200

KFAC*: damping: 7 . 978 934 · 10 -11 , exponential moving average: 8 . 950 193 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC: damping: 1 . 056 857 · 10 -13 , momentum: 7 . 160 131 · 10 -1 , exponential moving average: 9 . 622 372 · 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 ( U denotes a uniform, and LU a log-uniform distribution):

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

Adam: learning rate: LU ([ 5e-05 ; 5 · 10 -3 ])

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

KFAC*: damping: LU ([1 · 10 -12 ; 1 · 10 -6 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC: damping: LU ([1 · 10 -14 ; 1 · 10 -8 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

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

Setup Here, we consider three Poisson equations -∆ u ( x ) = f ( x ) with different right-hand sides and boundary conditions on the unit square x ∈ [0 , 1] d :

d = 5 with cosine sum right-hand side f ( x ) = π 2 ∑ d i =1 cos( πx i ) , boundary conditions u ( x ) = ∑ d i =1 cos( πx i ) for x ∈ ∂ [0 , 1] d , and known solution u ⋆ ( x ) = ∑ d i =1 cos( πx i ) . We assign each run a budget of 3 000 s.

d = 10 with zero right-hand side f ( x ) = 0 , harmonic mixed second order polynomial boundary conditions u ( x ) = ∑ d / 2 i =1 x 2 i -1 x 2 i for x ∈ ∂ [0 , 1] d , and known solution u ⋆ ( x ) = ∑ d / 2 i =1 x 2 i -1 x 2 i . We assign each run a budget of 6 000 s.

d = 100 with constant non-zero right-hand side f ( x ) = -2 d , square norm boundary conditions u ( x ) = ∥ x ∥ 2 2 for x ∈ ∂ [0 , 1] d , and known solution u ⋆ ( x ) = ∥ x ∥ 2 2 . We assign each run a budget of 10 000 s.

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

Best run details The runs shown in Figure A9 correspond to the following hyper-parameters:

5d Poisson equation, 5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116865

SGD: learning rate: 2 . 686 653 · 10 -4 , momentum: 9 . 878 243 · 10 -1 , N Ω : 606 , N ∂ Ω : 2 001 , batch sampling frequency: 1 570

Adam: learning rate: 6 . 111 767 · 10 -5 , N Ω : 534 , N ∂ Ω : 1 021 , batch sampling frequency: 8 220

-Hessian-free: curvature matrix: GGN, initial damping: 1 . 541 954 · 10 1 , constant damping: no, maximum CG iterations: 358 , N Ω : 1 084 , N ∂ Ω : 3 837 , batch sampling frequency: 230

LBFGS: learning rate: 1 . 749 124 · 10 -1 , history size: 339 , N Ω : 5 391 , N ∂ Ω : 4 768 , batch sampling frequency: 155

KFAC: damping: 4 . 251 462 · 10 -10 , momentum: 9 . 198 986 · 10 -1 , exponential moving average: 9 . 737 093 · 10 -1 , initialize Kronecker factors to identity: yes, N Ω : 4 690 , N ∂ Ω : 2 708 , batch sampling frequency: 2 369

KFAC*: damping: 2 . 240 865 · 10 -12 , exponential moving average: 8 . 522 194 · 10 -1 , initialize Kronecker factors to identity: yes, N Ω : 3 149 , N ∂ Ω : 3 801 , batch sampling frequency: 1 393

10d Poisson equation, 10 → 256 → 256 → 128 → 128 → 1 MLP with D = 118145

SGD: learning rate: 5 . 805 516 · 10 -2 , momentum: 9 . 715 522 · 10 -1 , N Ω : 537 , N ∂ Ω : 1 173 , batch sampling frequency: 1 083

Adam: learning rate: 1 . 337 679 · 10 -4 , N Ω : 115 , N ∂ Ω : 1 960 , batch sampling frequency: 4 975

-Hessian-free: curvature matrix: GGN, initial damping: 8 . 963 629 · 10 -1 , constant damping: no, maximum CG iterations: 143 , N Ω : 3 736 , N ∂ Ω : 961 , batch sampling frequency: 3

LBFGS: learning rate: 1 . 695 334 · 10 -1 , history size: 338 , N Ω : 342 , N ∂ Ω : 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, N Ω : 1 284 , N ∂ Ω : 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, N Ω : 1 090 , N ∂ Ω : 1 930 , batch sampling frequency: 2 454

100d Poisson equation, 100 → 768 → 768 → 512 → 512 → 1 MLP with D = 1325057

057

057

057

057

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

3

10

-

1

10

1

L

2

error

5d Poisson,

D

= 116

,

865

10

0

10

1

10

2

10

3

Time [s]

10

-

3

10

-

1

10d Poisson,

D

= 118

,

145

10

0

10

2

10

4

Time [s]

10

-

2

10

-

1

10

0

100d Poisson,

D

= 1

,

325

,

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

1

10

3

Loss

5d Poisson,

D

= 116

,

865

10

0

10

1

10

2

10

3

Time [s]

10

-

10

10

-

6

10

-

2

10d Poisson,

D

= 118

,

145

10

0

10

2

10

4

Time [s]

10

-

5

10

-

3

10

-

1

10

1

100d Poisson,

D

= 1

,

325

,

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

(b)

10

1

10

3

10

5

Iteration

10

-

3

10

-

1

10

1

L

2

error

5d Poisson,

D

= 116

,

865

10

1

10

3

10

5

Iteration

10

-

3

10

-

1

10d Poisson,

D

= 118

,

145

10

1

10

3

Iteration

10

-

2

10

-

1

10

0

100d Poisson,

D

= 1

,

325

,

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

10

1

10

3

10

5

Iteration

10

-

5

10

-

1

10

3

Loss

5d Poisson,

D

= 116

,

865

10

1

10

3

10

5

Iteration

10

-

10

10

-

6

10

-

2

10d Poisson,

D

= 118

,

145

10

1

10

3

Iteration

10

-

5

10

-

3

10

-

1

10

1

100d Poisson,

D

= 1

,

325

,

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

SGD: learning rate: 1 . 450 764 · 10 -3 , momentum: 9 . 747 671 · 10 -1 , N Ω : 177 , N ∂ Ω : 2 422 , batch sampling frequency: 519

Adam: learning rate: 7 . 894 685 · 10 -5 , N Ω : 100 , N ∂ Ω : 601 , batch sampling frequency: 19

-Hessian-free: curvature matrix: GGN, initial damping: 7 . 705 318 · 10 -4 , constant damping: no, maximum CG iterations: 263 , N Ω : 108 , N ∂ Ω : 2 372 , batch sampling frequency: 55

LBFGS: learning rate: 1 . 797 096 · 10 -1 , history size: 112 , N Ω : 2 115 , N ∂ Ω : 1 852 , 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, N Ω : 124 , N ∂ Ω : 2 332 , batch sampling frequency: 322

KFAC*: damping: 1 . 236 763 · 10 -7 , exponential moving average: 8 . 302 663 · 10 -1 , initialize Kronecker factors to identity: yes, N Ω : 175 , N ∂ Ω : 2 086 , batch sampling frequency: 16

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 ( U denotes a uniform, and LU a log-uniform distribution):

5d Poisson equation, 5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116865

Adam: learning rate: LU ([1 · 10 -6 ; 1]) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

SGD: learning rate: LU ([1 · 10 -6 ; 1]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: LU ([1 · 10 -15 ; 1]) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 1 , 2 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

KFAC: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

LBFGS: learning rate: LU ([1 · 10 -6 ; 1]) , history size: U ( { 5 , 6 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

KFAC*: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

SGD: learning rate: LU ([1 · 10 -6 ; 1]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 5 000 } )

10d Poisson equation, 10 → 256 → 256 → 128 → 128 → 1 MLP with D = 118145

Adam: learning rate: LU ([1 · 10 -6 ; 1]) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 5 000 } )

LBFGS: learning rate: LU ([1 · 10 -6 ; 1]) , history size: U ( { 5 , 6 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 5 000 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: LU ([1 · 10 -15 ; 1]) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 1 , 2 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 5 000 } )

KFAC: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 5 000 } )

100d Poisson equation, 100 → 768 → 768 → 512 → 512 → 1 MLP with D = 1325057

KFAC*: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 5 000 } )

SGD: learning rate: LU ([1 · 10 -6 ; 1]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 1 000 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: LU ([1 · 10 -15 ; 1]) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 1 , 2 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 1 000 } )

Adam: learning rate: LU ([1 · 10 -6 ; 1]) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 1 000 } )

LBFGS: learning rate: LU ([1 · 10 -6 ; 1]) , history size: U ( { 5 , 6 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 1 000 } )

KFAC*: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 1 000 } )

KFAC: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 5 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 2 500 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 1 000 } )

A.6 PINN Loss for the Heat Equation

Consider the ( ˜ d +1) -dimensional homogeneous heat equation

with spatial coordinates ˜ x ∈ Ω ⊆ R ˜ d and time coordinate t ∈ T ⊆ R where T is a time interval and κ > 0 denotes the heat conductivity. In this case, our neural network processes a ( d = ˜ d +1) -dimensional vector x = ( t, ˜ x ⊤ ) ⊤ ∈ R d and we can re-write the heat equation as

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

To make this more precise, consider the following example solution of the heat equation, which will be used later on as well. As initial conditions, we use u 0 (˜ x ) = u (0 , ˜ x ) = ∏ ˜ d i =1 sin( π ˜ x i ) for ˜ x ∈ Ω . For boundary conditions, we use g ( t, ˜ x ) = 0 for ( t, ˜ x ) ∈ T × ∂ Ω . The manufactured solution is

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

with x Ω n ∼ T × Ω , and x ∂ Ω n ∼ T × ∂ Ω , and x 0 n ∼ { 0 }× Ω . To fit this loss into our framework which assumes two loss terms, each of whose curvature is approximated with a Kronecker factor, we combine the initial value and boundary value conditions into a single term. Assuming N ∂ Ω = N 0 = N cond / 2 without loss of generality, we write

with domain inputs x Ω n ∼ T × Ω and targets y Ω n = 0 , boundary and initial condition targets y cond n = u ⋆ ( x cond n ) with initial inputs x cond n ∼ { 0 } × Ω for n = 1 , . . . , N cond / 2 and boundary inputs x cond n ∼ T × ∂ Ω for n = N cond / 2 +1 , . . . , N cond . This loss has the same structure as the PINN loss in Equation (1).

A.7 1+1d Heat Equation

Setup We consider a 1+1-dimensional heat equation ∂ t u ( t, x ) -κ ∆ x u ( t, x ) = 0 with κ = 1 / 4 on the unit square and unit time interval, x, t ∈ [0 , 1] × [0 , 1] . The equation has zero spatial boundary conditions and the initial values are given by u (0 , x ) = sin( πx ) for x ∈ [0 , 1] . We sample a single training batch of size N Ω = 900 , N ∂ Ω = 120 ( N ∂ Ω / 2 points for the initial value and spatial boundary conditions each) and evaluate the L 2 error on a separate set of 9 000 data points using the known solution u ⋆ ( t, x ) = exp( -π 2 t / 4 ) sin( πx ) . Each run is limited to 1 000 s. We compare three MLP architectures of increasing size, each of whose linear layers are Tanh-activated except for the final one: a shallow 2 → 64 → 1 MLP with D = 257 trainable parameters, a five layer 2 → 64 → 64 → 48 → 48 → 1 MLP with D = 9873 trainable parameters, and a five layer 2 → 256 → 256 → 128 → 128 → 1 MLP with D = 116097 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 A11 illustrates the learned solutions over training for all optimizers on the shallow MLP

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

2 → 64 → 1 MLP with D = 257

Adam: learning rate: 8 . 629 006 · 10 -4

SGD: learning rate: 1 . 752 752 · 10 -2 , momentum: 9 . 9 · 10 -1

-Hessian-free: curvature matrix: GGN, initial damping: 1 · 10 -4 , constant damping: no, maximum CG iterations: 350

LBFGS: learning rate: 1 · 10 -1 , history size: 125

-ENGD (full): damping: 1 · 10 -12 , exponential moving average: 9 · 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 -8 , momentum: 7 . 562 617 · 10 -1 , exponential moving average: 3 . 611 724 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC*: damping: 1 . 968 427 · 10 -9 , exponential moving average: 9 . 703 638 · 10 -1 , initialize Kronecker factors to identity: yes

2 → 64 → 64 → 48 → 48 → 1 MLP with D = 9873

Adam: learning rate: 2 . 551 515 · 10 -3

SGD: learning rate: 9 . 276 977 · 10 -2 , momentum: 9 . 9 · 10 -1

-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 -8 , exponential moving average: 6 · 10 -1 , initialize Gramian to identity: no

KFAC: damping: 3 . 169 186 · 10 -13 , momentum: 7 . 075 879 · 10 -1 , 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 → 1 MLP with D = 116097

Adam: learning rate: 6 . 716 485 · 10 -4

SGD: learning rate: 5 . 709 474 · 10 -2 , momentum: 9 . 9 · 10 -1

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

6

10

-

4

10

-

2

L

2

error

D

= 257

10

0

10

1

10

2

10

3

Time [s]

10

-

6

10

-

4

10

-

2

D

= 9

,

873

10

0

10

1

10

2

10

3

Time [s]

10

-

7

10

-

5

10

-

3

10

-

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

10

0

10

1

10

2

10

3

Time [s]

10

-

11

10

-

7

10

-

3

Loss

D

= 257

10

0

10

1

10

2

10

3

Time [s]

10

-

11

10

-

7

10

-

3

D

= 9

,

873

10

0

10

1

10

2

10

3

Time [s]

10

-

12

10

-

8

10

-

4

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

(b)

10

1

10

3

10

5

Iteration

10

-

6

10

-

4

10

-

2

L

2

error

D

= 257

10

1

10

3

10

5

Iteration

10

-

6

10

-

4

10

-

2

D

= 9

,

873

10

1

10

3

10

5

Iteration

10

-

7

10

-

5

10

-

3

10

-

1

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

10

1

10

3

10

5

Iteration

10

-

11

10

-

7

10

-

3

Loss

D

= 257

10

1

10

3

10

5

Iteration

10

-

11

10

-

7

10

-

3

D

= 9

,

873

10

1

10

3

10

5

Iteration

10

-

12

10

-

8

10

-

4

D

= 116

,

097

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

-Hessian-free: curvature matrix: GGN, initial damping: 1 · 10 -2 , constant damping: no, maximum CG iterations: 300

LBFGS: learning rate: 2 · 10 -1 , history size: 125

KFAC*: damping: 7 . 343 493 · 10 -11 , exponential moving average: 9 . 765 844 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC: damping: 2 . 576 488 · 10 -13 , momentum: 2 . 043 395 · 10 -2 , exponential moving average: 9 . 727 829 · 10 -1 , initialize Kronecker factors to identity: yes

Search space details The runs shown in Figure A10 were determined to be the best via a search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search ( U denotes a uniform, and LU a log-uniform distribution):

2 → 64 → 1 MLP with D = 257 - SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -1 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) - Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

Optimizer

First step

0.1% trained

1%trained

10% trained

True solution

SGD

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

6

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

277

, Loss:

2

.

5

·

10

-

2

,

L

2

loss:

1

.

0

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

3000

, Loss:

6

.

1

·

10

-

4

,

L

2

loss:

1

.

4

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

29540

, Loss:

2

.

2

·

10

-

5

,

L

2

loss:

1

.

4

·

10

-

3

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

6

·

10

-

1

Adam

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

7

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

277

, Loss:

5

.

5

·

10

-

2

,

L

2

loss:

1

.

7

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

2727

, Loss:

7

.

1

·

10

-

4

,

L

2

loss:

1

.

4

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

26855

, Loss:

6

.

8

·

10

-

6

,

L

2

loss:

7

.

3

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

7

·

10

-

1

LBFGS

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

1

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

55

, Loss:

1

.

2

·

10

-

5

,

L

2

loss:

1

.

0

·

10

-

3

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

594

, Loss:

2

.

8

·

10

-

6

,

L

2

loss:

2

.

7

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

5845

, Loss:

2

.

8

·

10

-

6

,

L

2

loss:

2

.

7

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

1

·

10

-

1

Hessian-free

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

3

.

4

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

1

, Loss:

1

.

2

·

10

-

1

,

L

2

loss:

3

.

5

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

11

, Loss:

7

.

5

·

10

-

3

,

L

2

loss:

8

.

6

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

107

, Loss:

4

.

6

·

10

-

8

,

L

2

loss:

7

.

8

·

10

-

5

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

3

.

4

·

10

-

1

ENGD (full)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

3

.

5

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

4

, Loss:

1

.

1

·

10

-

1

,

L

2

loss:

3

.

2

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

42

, Loss:

2

.

3

·

10

-

3

,

L

2

loss:

2

.

1

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

446

, Loss:

4

.

9

·

10

-

12

,

L

2

loss:

6

.

6

·

10

-

7

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

3

.

5

·

10

-

1

ENGD

(layer-wise)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

6

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

4

, Loss:

1

.

0

·

10

-

1

,

L

2

loss:

2

.

7

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

42

, Loss:

4

.

2

·

10

-

2

,

L

2

loss:

1

.

5

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

446

, Loss:

1

.

8

·

10

-

4

,

L

2

loss:

1

.

1

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

6

·

10

-

1

KFAC

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

1

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

14

, Loss:

1

.

8

·

10

-

2

,

L

2

loss:

7

.

5

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

143

, Loss:

2

.

5

·

10

-

3

,

L

2

loss:

2

.

5

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

1400

, Loss:

4

.

7

·

10

-

7

,

L

2

loss:

3

.

2

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

1

·

10

-

1

KFAC*

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

7

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

Normalized true solution

Step:

34

, Loss:

1

.

4

·

10

-

2

,

L

2

loss:

7

.

6

·

10

-

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

Normalized true solution

Step:

335

, Loss:

1

.

5

·

10

-

5

,

L

2

loss:

7

.

9

·

10

-

4

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

Normalized true solution

Step:

3299

, Loss:

5

.

3

·

10

-

9

,

L

2

loss:

2

.

3

·

10

-

5

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

t

Normalized learned solution

Normalized true solution

Step:

0

, Loss:

1

.

4

·

10

-

1

,

L

2

loss:

2

.

7

·

10

-

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

0

.

0

0

.

0

0

.

2

0

.

0

2

.

4

0

.

0

4

.

6

0

.

0

6

.

8

0

.

1

8

.

0

1

.

0

0

.

0

0

.

0

0

.

2

0

.

0

2

.

4

0

.

0

4

.

6

0

.

0

6

.

8

0

.

1

8

.

0

1

.

0

0

.

0

0

.

0

0

.

2

0

0

.

.

0

2

0

.

4

0

0

.

.

0

4

2

.

6

0

0

.

.

0

6

4

.

8

0

0

.

.

1

8

6

.

0

1

0

.

.

0

8

0

.

0

1

.

0

0

0

.

.

2

0

0

0

.

.

4

2

0

0

.

.

6

4

0

0

.

.

8

6

1

0

.

.

0

8

1

.

0

x

x

x

x

x

x

x

x

x

x

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 100 , 1 , 1 · 10 -2 , 1 · 10 -4 , 1 · 10 -6 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

-ENGD (full): damping: U ( { 1 · 10 -6 , 1 · 10 -8 , 1 · 10 -10 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 75 , 100 , 125 , 150 , 175 , 200 , 225 , 250 } )

-ENGD(layer-wise): damping: U ( { 1 · 10 -4 , 1 · 10 -6 , 1 · 10 -8 , 1 · 10 -10 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC*: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -1 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } )

2 → 64 → 64 → 48 → 48 → 1 MLP with D = 9873

Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 75 , 100 , 125 , 150 , 175 , 200 , 225 , 250 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 100 , 1 , 1 · 10 -2 , 1 · 10 -4 , 1 · 10 -6 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

-ENGD (full): damping: U ( { 1 · 10 -6 , 1 · 10 -8 , 1 · 10 -10 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

-ENGD(layer-wise): damping: U ( { 1 · 10 -4 , 1 · 10 -6 , 1 · 10 -8 , 1 · 10 -10 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC*: damping: LU ([1 · 10 -15 ; 1 · 10 -9 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -1 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } )

2 → 256 → 256 → 128 → 128 → 1 MLP with D = 116097

Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 75 , 100 , 125 , 150 , 175 , 200 , 225 , 250 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 100 , 1 , 1 · 10 -2 , 1 · 10 -4 , 1 · 10 -6 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

KFAC: damping: LU ([1 · 10 -14 ; 1 · 10 -7 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC*: damping: LU ([1 · 10 -14 ; 1 · 10 -6 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

A.8 4+1d Heat Equation

Setup We consider a 4+1-dimensional heat equation ∂ t u ( t, x ) -κ ∆ x u ( t, x ) = 0 with κ = 1 / 4 on the four-dimensional unit square and unit time interval, x , t ∈ [0 , 1] 4 × [0 , 1] . The equation has spatial boundary conditions u ( t, x ) = exp( -t ) ∑ 4 i =1 sin(2 x i ) for t, x ∈ [0 , 1] × ∂ [0 , 1] 4 throughout time, and initial value conditions u (0 , x ) = ∑ 4 i =1 sin(2 x i ) for x ∈ [0 , 1] 4 . We sample training batches of size N Ω = 3000 , N ∂ Ω = 500 ( N ∂ Ω / 2 points for the initial value and spatial boundary conditions each) and evaluate the L 2 error on a separate set of 30 000 data points using the known solution u ⋆ ( t, x ) = exp( -t ) ∑ 4 i =1 sin(2 x i ) . All optimizers except for KFAC sample a new training batch each iteration. KFAC only re-samples every 100 iterations because we noticed significant improvement with multiple iterations on a fixed batch. To make sure that this does not lead to an unfair advantage of KFAC, we conduct an additional experiment where we also tune the batch sampling frequency, as well as other hyper-parameters; see §A.10. The results presented in this section are consistent with this additional experiment (compare the rightmost column of Figure A12

and Figure A14). Each run is limited to 3000 s. We compare three MLP architectures of increasing size, each of whose linear layers are Tanh-activated except for the final one: a shallow 5 → 64 → 1 MLP with D = 449 trainable weights, a five layer 5 → 64 → 64 → 48 → 48 → 1 MLP with D = 10065 trainable weights, and a five layer 5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116864 trainable weights. For the biggest architecture, full and layer-wise ENGD lead to out-of-memory errors and are thus not tested. Figure A12 visualizes the results.

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

L

2

error

D

= 449

10

0

10

1

10

2

10

3

Time [s]

10

-

4

10

-

2

10

0

D

= 10

,

065

10

0

10

1

10

2

10

3

Time [s]

10

-

4

10

-

2

10

0

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

9

10

-

5

10

-

1

Loss

D

= 449

10

0

10

1

10

2

10

3

Time [s]

10

-

7

10

-

3

10

1

D

= 10

,

065

10

0

10

1

10

2

10

3

Time [s]

10

-

8

10

-

5

10

-

2

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

(b)

10

1

10

3

10

5

Iteration

10

-

5

10

-

3

10

-

1

L

2

error

D

= 449

10

1

10

3

10

5

Iteration

10

-

4

10

-

2

10

0

D

= 10

,

065

10

1

10

3

Iteration

10

-

4

10

-

2

10

0

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

10

1

10

3

10

5

Iteration

10

-

9

10

-

5

10

-

1

Loss

D

= 449

10

1

10

3

10

5

Iteration

10

-

7

10

-

3

10

1

D

= 10

,

065

10

1

10

3

Iteration

10

-

8

10

-

5

10

-

2

D

= 116

,

865

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

KFAC

KFAC*

Search space details The runs shown in Figure A12 were determined to be the best via a search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search ( U denotes a uniform, and LU a log-uniform distribution):

5 → 64 → 1 MLP with D = 449

SGD: learning rate: 7 . 737 742 · 10 - 3 , momentum: 9 · 10 - 1

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: 1 · 10 - 10 , exponential moving average: 6 · 10 - 1 , initialize Gramian to identity: yes

-ENGD (layer-wise): damping: 1 · 10 -6 , exponential moving average: 0 , initialize Gramian to identity: yes

KFAC: damping: 1 . 000 288 · 10 -9 , momentum: 9 . 474 108 · 10 -1 , exponential moving average: 7 . 783 519 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC*: damping: 2 . 965 060 · 10 -8 , exponential moving average: 9 . 574 717 · 10 -1 , initialize Kronecker factors to identity: yes

5 → 64 → 64 → 48 → 48 → 1 MLP with D = 10065

Adam: learning rate: 7 . 801 748 · 10 -4

SGD: learning rate: 9 . 357 973 · 10 -3 , momentum: 9 · 10 -1

-Hessian-free: curvature matrix: GGN, initial damping: 5 · 10 -3 , constant damping: no, maximum CG iterations: 400

LBFGS: learning rate: 1 · 10 -1 , history size: 225

-ENGD (full): damping: 1 · 10 -8 , exponential moving average: 0 , initialize Gramian to identity: yes

-ENGD(layer-wise): damping: 1 · 10 -6 , exponential moving average: 3 · 10 -1 , initialize Gramian to identity: no

KFAC: damping: 4 . 143 385 · 10 -14 , momentum: 7 . 660 303 · 10 -1 , exponential moving average: 9 . 821 414 · 10 -1 , initialize Kronecker factors to 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 → 1 MLP with D = 116865

Adam: learning rate: 5 . 266 284 · 10 -4

SGD: learning rate: 7 . 192 473 · 10 -3 , momentum: 9 · 10 -1

-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: 3 . 405 440 · 10 -14 , exponential moving average: 8 . 445 471 · 10 -1 , initialize Kronecker factors to identity: yes

KFAC: damping: 8 . 581 322 · 10 -13 , momentum: 8 . 501 747 · 10 -1 , exponential moving average: 9 . 803 115 · 10 -1 , initialize Kronecker factors to identity: yes

Search space details The runs shown in Figure A12 were determined to be the best via a search with approximately 50 runs on the following search spaces which were obtained by refining an initially wider search ( U denotes a uniform, and LU a log-uniform distribution):

5 → 64 → 1 MLP with D = 449

Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

-ENGD (full): damping: U ( { 1 · 10 -8 , 1 · 10 -9 , 1 · 10 -10 , 1 · 10 -11 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

-ENGD (layer-wise): damping: U ( { 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 , 1 · 10 -5 , 1 · 10 -6 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC*: damping: LU ([1 · 10 -13 ; 1 · 10 -7 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC: damping: LU ([1 · 10 -12 ; 1 · 10 -6 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

5 → 64 → 64 → 48 → 48 → 1 MLP with D = 10065

Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

-ENGD (full): damping: U ( { 1 · 10 -8 , 1 · 10 -9 , 1 · 10 -10 , 1 · 10 -11 , 1 · 10 -12 , 0 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

-ENGD (layer-wise): damping: U ( { 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 , 1 · 10 -5 , 1 · 10 -6 } ) , exponential moving average: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 , 9 . 9 · 10 -1 } ) , initialize Gramian to identity: U ( { no , yes } )

KFAC*: damping: LU ([1 · 10 -14 ; 1 · 10 -8 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

KFAC: damping: LU ([1 · 10 -14 ; 1 · 10 -8 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116865

Adam: learning rate: LU ([5 · 10 -4 ; 1 · 10 -1 ])

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

-Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

KFAC: damping: LU ([1 · 10 -14 ; 1 · 10 -8 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: yes

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

KFAC*: damping: LU ([1 · 10 -14 ; 1 · 10 -8 ]) , exponential moving average: U ([5 · 10 -1 ; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to 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 D = 10065 parameters from Figure 2's middle panel which is bigger than the two-layer toy model, while still allowing to run ENGD. Using the same hyper-parameters, 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 as in §A.8 and use the 5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116865 . We tune all optimizer hyper-parameters as described in §A.1 and also tune the batch sizes N Ω , N ∂ Ω , as well as their re-sampling frequencies. Figure A14 summarizes the results.

Best run details The runs shown in Figure A14 correspond to the following hyper-parameters:

SGD: learning rate: 1 . 614 965 · 10 -2 , momentum: 9 . 899 167 · 10 -1 , N Ω : 527 , N ∂ Ω : 2 157 , batch sampling frequency: 543

Adam: learning rate: 2 . 583 569 · 10 -4 , N Ω : 472 , N ∂ Ω : 3 018 , batch sampling frequency: 177

4d Heat (mlp-tanh-64-64-48-48,

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

10

1

L

2

error

SGD

Adam

Hessian-free

LBFGS

ENGD (full)

ENGD (layer-wise)

ENGD (diagonal)

KFAC

KFAC*

D

= 10065

)

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

5

10

-

3

10

-

1

L

2

error

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

11

10

-

8

10

-

5

10

-

2

10

1

Loss

(b)

10

0

10

1

10

2

10

3

10

4

10

5

Iteration

10

-

5

10

-

3

10

-

1

L

2

error

10

0

10

1

10

2

10

3

10

4

10

5

Iteration

10

-

11

10

-

8

10

-

5

10

-

2

10

1

Loss

Hessian-free: curvature matrix: GGN, initial damping: 5 . 077 634 · 10 -4 , constant damping: yes, maximum CG iterations: 163 , N Ω : 1 172 , N ∂ Ω : 1 637 , batch sampling frequency: 5 440

LBFGS: learning rate: 1 . 029 194 · 10 -1 , history size: 488 , N Ω : 582 , N ∂ Ω : 2 038 , batch sampling frequency: 315

KFAC: damping: 8 . 435 180 · 10 -14 , momentum: 9 . 718 645 · 10 -1 , exponential moving average: 9 . 800 744 · 10 -1 , initialize Kronecker factors to identity: yes, N Ω : 2 525 , N ∂ Ω : 2 663 , batch sampling frequency: 7 916

KFAC*: damping: 8 . 837 871 · 10 -15 , exponential moving average: 9 . 887 596 · 10 -1 , initialize Kronecker factors to identity: yes, N Ω : 2 563 , N ∂ Ω : 2 873 , batch sampling frequency: 9 647

Search space details The runs shown in Figure A14 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 ( U denotes a uniform, and LU a log-uniform distribution):

SGD: learning rate: LU ([1 · 10 -6 ; 1]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: LU ([1 · 10 -15 ; 1]) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 1 , 2 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

Adam: learning rate: LU ([1 · 10 -6 ; 1]) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

LBFGS: learning rate: LU ([1 · 10 -6 ; 1]) , history size: U ( { 5 , 6 , . . . , 500 } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

KFAC*: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

KFAC: damping: LU ([1 · 10 -15 ; 1 · 10 -2 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([0; 9 . 9 · 10 -1 ]) , initialize Kronecker factors to identity: U ( { no , yes } ) , N Ω : U ( { 100 , 101 , . . . , 10 000 } ) , N ∂ Ω : U ( { 50 , 51 , . . . , 5 000 } ) , batch sampling frequency: U ( { 0 , 1 , . . . , 10 000 } )

A.11 9+1-d Logarithmic Fokker-Planck Equation with Random Search

For a given drift µ : [0 , 1] × R d → R d and diffusivity σ : [0 , 1] → R d × d the Fokker-Planck equation with initial probability density p 0 is given by

which is posed on [0 , 1] × R d . Note that p ( t, · ) is a probability density on R d for all t ∈ [0 , 1] . We transform the above equation into logarithmic space via q = log( p ) . Then q solves

For the concrete example of the main text, we set µ ( t, x ) = -1 2 x and σ = √ 2 I ∈ R d × d . We consider a 9+1 dimensional Fokker-Planck equation in logarithmic space and replace the unbounded domain by [0 , 1] × [ -5 , 5] d . Precisely, we aim to solve the equation

where d = 9 , t ∈ [0 , 1] and x ∈ [ -5 , 5] . The solution q ∗ = log( p ∗ ) is given as p ∗ ( t, x ) ∼ N (0 , exp( -t ) I + (1 -exp( -t ))2 I ) . The PINN loss includes the PDE residual and the initial conditions. We model the solution with a medium sized tanh-activated MLP with D = 118145 and the layer structure 10 → 256 → 256 → 128 → 128 → 1 and use batch sizes of N Ω = 3000 , N ∂ Ω = 1000 . Each run is assigned a budget of 6 000 s. Figure A15 visualizes the results.

(a)

10

0

10

1

10

2

10

3

Time [s]

10

-

1

10

0

10

1

L

2

error

SGD

Adam

Hessian-free

LBFGS

KFAC

KFAC*

10

0

10

1

10

2

10

3

Time [s]

10

-

2

10

0

10

2

Loss

(b)

10

0

10

1

10

2

10

3

10

4

Iteration

10

-

1

10

0

10

1

L

2

error

10

0

10

1

10

2

10

3

10

4

Iteration

10

-

2

10

0

10

2

Loss

Search space details The runs shown in Figure A15 were determined to be the best via a random search on the following search spaces which each optimizer given approximately the same total computational time ( U denotes a uniform, and LU a log-uniform distribution):

SGD: learning rate: LU ([1 · 10 -3 ; 1 · 10 -2 ]) , momentum: U ( { 0 , 3 · 10 -1 , 6 · 10 -1 , 9 · 10 -1 } )

Hessian-free: curvature matrix: U ( { GGN , Hessian } ) , initial damping: U ( { 1 , 1 · 10 -1 , 1 · 10 -2 , 1 · 10 -3 , 1 · 10 -4 } ) , constant damping: U ( { no , yes } ) , maximum CG iterations: U ( { 50 , 250 } )

Adam: learning rate: LU ([ 5e-05 ; 5 · 10 -3 ])

LBFGS: learning rate: U ( { 5 · 10 -1 , 2 · 10 -1 , 1 · 10 -1 , 5 · 10 -2 , 2 · 10 -2 , 1 · 10 -2 } ) , history size: U ( { 50 , 75 , 100 , 125 , 150 , 175 , 200 , 225 } )

KFAC*: damping: LU ([1 · 10 -11 ; 1 · 10 -5 ]) , exponential moving average: U ([9 . 99 · 10 -1 ; 9 . 999 · 10 -1 ]) , initialize Kronecker factors to identity: yes We found that KFAC* requires very large exponential moving averages to work well.

KFAC: damping: LU ([1 · 10 -11 ; 1 · 10 -5 ]) , momentum: U ([0; 9 . 9 · 10 -1 ]) , exponential moving average: U ([9 . 99 · 10 -1 ; 9 . 999 · 10 -1 ]) , initialize Kronecker factors to identity: yes

B Pseudo-Code: KFAC for the Poisson Equation

Algorithm 1 KFAC for the Poisson equation.

Require:

boundary data { ( x b n , y b n ) } ∂ Ω n =1 exponential moving average β , momentum µ , Damping λ , number of steps T

0) Initialization

1) Compute the interior loss and update its approximate curvature

( Z (0) n . . . , Z ( L ) n , ∆ u n ) ← ∆ u θ t ( x n ) n = 1 , . . . , N Ω ▷ Forward Laplacian wit intermediates Compute layer output gradients g ( l ) n,s := ∂ ∆ u n / ∂ Z ( l ) n,s with autodiff in one backward pass

(

g

(1)

n,s

, . . . , for all

l

g

(

L

)

n,s

= 1

)

, . . . , L

grad do

= 1

▷

end for

2) Compute the boundary loss and update its approximate curvature

for all

l

= 1

, . . . , L

do

▷

Update Kronecker factors of the boundary loss

end for

3) Update the preconditioner (use inverse of Kronecker sum trick)

end for

4) Compute the gradient using autodiff, precondition the gradient

▷ Gradient with autodiff

▷ Precondition gradient

▷ Proposed update direction

Add momentum from previous update

▷ Parameter update

end for

5) Given the direction ˆ δ (1) t , . . . , ˆ δ ( L ) t , choose learning rate α by line search & update

end for end for

return Trained parameters θ T

(∆

u

,

(

Z

(1)

n,s

, . . . ,

Z

(

L

)

n,s

))

n

, . . . , N

Ω

,

s

= 1

, . . . , S

:=

d

+2

Update Kronecker factors of the interior loss

▷ Compute interior loss

n

▷ Initialize Kronecker factors

▷ Compute boundary loss

C Taylor-Mode Automatic Differentiation & Forward Laplacian

PINN losses involve differential 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 a Kroneckerfactorized approximation of the Gramian, which consists of the Laplacian's gradient, we need to know how a weight matrix enters its computation. Here, we describe how the weight matrix of a linear layer inside a feed-forward net enters the Laplacian's computation when using the forward Laplacian framework. We start by connecting the forward Laplacian framework to Taylor-mode automatic differentiation [18, 3], 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 to forward-propagate Taylor coefficients, i.e. directional derivatives, through the computation graph. We provide a brief summary based on its description in [3].

Taylor series and directional derivatives Consider a function f : R d → R and its K -th order Taylor expansion at a point x ∈ R d along a direction α v ∈ R d with α ∈ R ,

We can unify this expression by introducing the K -th order directional derivative of f at x along v ,

This simplifies the uni-directional Taylor expansion to

where we have used the notation ⊗ k v to indicate k copies of v , and introduced the k -th order Taylor coefficient w f k ∈ R of f . This generalizes to vector-valued functions: If f 's output was vector-valued, say f ( x ) ∈ R c , we would have Taylor-expanded each component individually and grouped coefficients of same order into vectors w f k ∈ R c such that [ w f k ] i is the k -th order Taylor coefficient of the i th component of f .

̸

A note on generality: In this introduction to Taylor-mode, we limit the discussion to the computation of higher-order derivatives along a single direction v , i.e. ∂ K f ( x )[ v , . . . , v ] . This is limited though, e.g. if we set K = 2 then we can compute ∂ 2 f ( x )[ v , v ] = v ⊤ ( ∂ 2 f ( x ) / ∂ x 2 ) v . We can set v = e i to the i -th standard basis vector to compute the i -th diagonal element of the Hessian. But we cannot evaluate off-diagonal elements, as this would require multi-directional derivatives, like ∂ 2 f ( x )[ e i , e j = i ] . A more general description of Taylor-mode for multi-directional Taylor series along M directions, ˆ f ( α 1 , . . . , α M ) = f ( x + α 1 v 1 + · · · + α M v M ) , which require more general directional derivatives ∂ K f ( x )[ v 1 , . . . , v K ] (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 §C.3.

Composition rule Next, we consider the case where f = g ◦ h is a composition of two functions. Starting from the Taylor coefficients w h 0 , . . . w h K of ˆ h ( α ) = h ( x + α v ) , the Taylor coefficients w f 0 , . . . , w f K of ˆ f ( α ) = f ( x + α v ) follow from Faà di Bruno's formula [18, 3]:

In the above, part( k ) is the set of all integer partitionings of k ; a set of sets. | σ | denotes the length of a set σ ∈ part( k ) , n i is the count of integer i in σ , and w h 0 = h ( x ) .

Second-order Taylor-mode Our goal is the computation of second-order derivatives of f w.r.t. x . So let's work out Equation (C20) up to order K = 2 . The zeroth and first order are simply the forward pass and the forward-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

We can also express w f 1 , w f 2 in terms of Jacobian- and Hessian-vector products of g ,

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 is to compute the Laplacian of f : R d → R c (in practise, c = 1 ),

where e i is the i -th standard basis vector, [ f ( x )] j is the j -th component of f ( x ) , and we have introduced the second-order Taylor coefficients w f 2 ,i of f along e i . The Laplacian requires computing, then summing, the second-order Taylor coefficients of d Taylor approximations { f ( x + e i ) } i =1 ,...,d .

Naive approach We can use Taylor-mode differentiation to compute all these components in 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),

Forward Laplacian framework Computing the Laplacian via Equation (C24) first computes, then sums, the diagonal second-order derivatives { w f 2 ,i } i =1 ,...,d . Note that we can pull the sum

inside the forward propagation scheme, specifically Equation (C24c), and push-forward the summed second-order coefficients. This simplifies Equation (C24) to

Equation (C25) is the forward Laplacian framework from Li et al. [29] for computing the Laplacian of a neural network. Here, we have derived it 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 : R c → R c acts element-wise via σ : R → R . We will write σ ( · ) , σ ′ ( · ) , σ ′′ ( · ) to indicate element-wise application of σ , its first derivative σ ′ , and second derivative σ ′′ to all elements of · . Further, let ⊙ denote element-wise multiplication, and ( · ) ⊙ 2 element-wise squaring. With that, we can write the Jacobian as J h ( x ) g ( x ) = diag( σ ( h ( x ))) where diag( · ) embeds a vector · into the diagonal of a matrix. The Hessian of component i is ∂ 2 [ g ( h ( x ))] i / ∂h ( x ) 2 = [ σ ′′ ( h ( x ))] i e i e ⊤ i . 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:

C.2.2 Forward Laplacian for Linear Layers

Now, let g : R D in → R D out be a linear layer with weight matrix W ∈ R D out × D in and bias vector b ∈ R D out . Its Jacobian is J h ( x ) ( W h ( x ) + b ) = W and the second-order derivative is zero. Hence, Equation (C25) for linear layers becomes

We can summarize Equation (C27) in a single equation by grouping all quantities that are multiplied by W into a single matrix, and appending a single row of ones or zeros to account for the bias:

Equation (C28) shows that the weight matrix ˜ W ( l ) = ( W ( l ) b ( l ) ) of a linear layer f ( l ) inside a neural network f ( L ) ◦ . . . ◦ f (1) is applied to a matrix T ( l -1) ∈ R D in × ( d +2) during the computation of the net's prediction and Laplacian via the forward Laplacian framework and yields another matrix T ( l ) ∈ R D out × ( d +2) .

or, in compact form,

C.3 Generalization of the Forward Laplacian to Weighted Sums of Second Derivatives

The Laplacian is of the form ∆ x f = ∑ i ∂ 2 f ( x )[ e i , e i ] 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 ∑ i,j c i,j ∂ 2 f ( x )[ e i , e j ] , which contain the Laplacian for c i,j = δ i,j .

As mentioned in §C.1, this requires a generalization of Taylor-mode which computes derivatives of the form ∂ K f ( x )[ v , . . . , v ] , where the directions v must be identical. We start with the formulation in [25] which expresses the K -th multi-directional derivative of a function f = g ◦ h through the composites' derivatives (all functions can be vector-to-vector)

Here, part( { 1 , . . . , K } ) denotes the set of all set partitions of { 1 , . . . , K } ( σ is a set of sets). E.g.,

To make this more concrete, let's consider Equation (C29) for first- and second-order derivatives,

From Equation (C30), we can see that if we want to compute a weighted sum of second-order derivatives ∑ i,j c i,j ∂ 2 f ( x )[ v i , v j ] , we can pull the sum inside the second equation,

Hence, we can propagate the collapsed second-order derivatives, together with all first-order derivatives along v 1 , v 2 , . . . . The only difference to the forward Laplacian is 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 a linear layer g : h ( x ) ↦→ W h ( x ) + b , we have ∂ m> 1 g ( h ( x ))[ v 1 , . . . , v m ] = 0 , and thus

The last equation is because only the set partition { 1 , . . . , K } contributes to Equation (C29).

For elementwise activations g : h ( x ) ↦→ σ ( h ( x )) with σ : R → R applied component-wise, we have the structured derivative tensor [ ∂ m g ( h ( x ))] i 1 ,...,i m = ∂ m σ ( h ( x ) i 1 ) δ i 1 ,...,i m and multi-directional derivative ∂ K g ( h ( x ))[ v 1 , . . . , v K ] = ∂ K σ ( x ) ⊙ v 1 ⊙··· ⊙ v K . Equation (C30) becomes

As shown in Equation (C30b), for both Equations (C32) and (C33), we can pull the summation inside the propagation scheme. Specifically, to compute ∑ i,j c i,j ∂ 2 f ( x )[ e i , e j ] , we have for linear layers

and for activation layers

(the summed second-order derivatives that are forward-propagated are highlighted). This propagation reduces back to the forward Laplacian Equations (C26) and (C27) when we set c i,j = δ i,j . In contrast to other attempts to compute such a 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 an off-the shelve solution. We consider two Laplacian implementations:

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.

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 D Ω → 768 → 768 → 512 → 512 → 1 MLP with tanh-activations from Figure 3), then measure run time and peak memory of computing the net's Laplacian on a mini-batch of size N = 1024 with varying values of D Ω . To reduce measurement noise, we repeat each run over five independent Python sessions and report the smallest value (using the same GPU as in all other experiments, an NVIDIA RTX 6000 with 24 GiB memory).

We observe that the forward Laplacian is roughly twice as fast as the functorch Laplacian, and that it uses significantly less memory for large input dimensions, up to only one half when D Ω = 100 . We visualized both tables using more values for D Ω , see Figure C16. In the shown regime, we find that the MLP's increasing cost in D Ω (due to the growing first layer) is negligible as we observe linear scaling in both memory and run time. For extremely large D Ω , it would eventually become quadratic.

D Backpropagation Perspective of the Laplacian

Here, we 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 as simple weight-sharing. This complicates defining a Kroneckerfactored approximation for the Gramian without introducing new approximations that are different from Eschenhagen et al. [17], rendering the Taylor-mode perspective advantageous.

We start by deriving the Laplacian ∆ u := Tr( ∇ 2 x u ) of a feed-forward NN (see §2.1), assuming a single data point for simplicity (see §D.1) and abbreviating u θ as u . The goal is to make the Laplacian's dependence w.r.t. a weight W ( i ) in one layer of the network explicit. Then, we can write

down the Jacobian J W ( i ) ∆ u (see §D.2) which is required for the Gramian in Equation (2) (see §D.3). We do this based on the concept of Hessian backpropagation [10, HBP,] which yields a recursion for the Hessian ∇ 2 x u . The Laplacian follows by taking the trace of the latter. Finally, we use the chain rule express the Laplacian's Jacobian J W ( i ) ∆ u in terms of W ( i ) 's children in 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 to compute Hessians w.r.t. nodes in a graph ( z ( i ) or θ ( i ) ). We call this recursive procedure Hessian backpropagation [10].

Gradient backpropagation As a warm-up, let's recall how to compute the gradient ∇ θ u = ( ∇ θ (1) u, . . . , ∇ θ ( L ) u ) . We start by setting ∇ z ( L ) u = ∇ u u = 1 (assuming u is scalar for simplicity), then backpropagate the error via VJPs according to the recursion

for i = L, . . . , 1 . This yields the gradients of u w.r.t. all intermediate representations and parameters.

Hessian backpropagation Just like gradient backpropagation, we can derive a recursive scheme for the Hessian. Recall the Hessian chain rule

where g i denotes the individual components of g , see [53]. The recursion for computing Hessians of u w.r.t. intermediate representations and parameters starts by initializing the recursion with ∇ 2 z ( L ) u = ∇ 2 u u = 0 , and then backpropagating according to (see Dangel et al. [10] for details)

for i = L, . . . , 1 . The first term takes the incoming Hessian (w.r.t. a layer's output) and sandwiches it between the layer's Jacobian. It can be seen as backpropagating curvature from downstream layers. The second term adds in curvature introduced by the current layer. It is only non-zero if 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 z ( i ) u, ∇ 2 θ ( i ) u . In Figure D17 we depict the dependencies of intermediate gradients and Hessians for computing ∇ 2 x u = ∇ 2 z (0) u :

Parameters

θ

(1)

. . .

θ

(

i

-

1)

θ

(

i

)

. . .

θ

(

L

)

Forward

x

z

(1)

. . .

z

(

i

-

1)

z

(

i

)

. . .

u

Backward

∇

x

u

∇

z

(1)

u

. . .

∇

z

(

i

-

1)

u

∇

z

(

i

)

u

. . .

∇

u

u

Hess. backward

∆

u

∇

2

x

u

∇

2

z

(1)

u

. . .

∇

2

z

(

i

-

1)

u

∇

2

z

(

i

)

u

. . .

∇

2

u

u

∇ z ( i -1) u depends on ∇ z ( i ) u due to the recursion in Equation (D35), and on z ( i -1) , θ ( i ) due to the Jacobian J z ( i -1) z ( i ) in the gradient backpropagation Equation (D35).

∇ 2 z ( i -1) u depends on ∇ 2 z ( i ) u and ∇ z ( i ) u due to the recursion in Equation (D37), and on z ( i -1) , θ ( i ) due to the Jacobian J z ( i -1) z ( i ) and Hessian ∇ 2 z ( i -1) [ z ( i ) ] k in the Hessian backpropagation Equation (D35).

The Laplacian ∆ u follows by taking the trace of ∇ 2 x u 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-wise nonlinear activations.

Hessian backpropagation through nonlinear layers We mostly consider nonlinear layers without trainable parameters and consist of a componentwise nonlinearity z ↦→ σ ( z ) for some σ : R → R . The Jacobian of such a nonlinear layer is given by J z ( i -1) z ( i ) = diag( σ ′ ( z ( i -1) )) and the Hessian terms are given by ∇ 2 z ( i -1) [ z ( i ) ] k = σ ′′ ( z ( i -1) k ) e k e ⊤ k where e k is the unit vector along coordinate k . With these two identities we can backpropogate the input Hessian through such layers via

Hessian backpropagation through a linear layer To de-clutter the dependency graph of Figure D17, we will now consider the dependency of ∆ u w.r.t. the weight of a single layer. We assume this layer i to be a linear layer with parameters W ( i ) such that θ ( i ) = vec( W ( i ) ) ,

For this layer, the second terms in Equation (D37) disappears because the local Hessians are zero, that is ∇ 2 z ( i -1) [ z ( i ) ] k = 0 and ∇ 2 W ( i ) [ z ( i ) ] k = 0 . Also, the Jacobians are J W ( i ) z ( i ) = z ( i -1) ⊤ ⊗ I and J z ( i -1) z ( i ) = W ( i ) and hence only depend on one of the two layer inputs. This simplifies the computation graph. Figure D18 shows the dependencies of W ( i ) on the Laplacian, highlighting its three direct children,

. . .

W

(

i

)

. . .

. . .

z

(

i

-

1)

z

(

i

)

. . .

. . .

∇

z

(

i

-

1)

u

∇

z

(

i

)

u

. . .

∆

u

. . .

∇

2

z

(

i

-

1)

u

∇

2

z

(

i

)

u

. . .

Figure D18: Direct dependencies of a linear layer's weight matrix W ( i ) in the Laplacian's computation graph. There are three direct children: (i) the layer's output from the forward pass, (ii) the Laplacian's gradient w.r.t. the layer's input from the gradient backpropagation, and (iii) the Laplacian's Hessian w.r.t. the layer's input from the Hessian backpropagation. The Jacobians J W ( i ) ∆ u required for the Gramian are the vector-Jacobian products accumulated over those children.

D.2 Parameter Jacobian of the Backward Laplacian

Recall that the entries of the Gramian are composed from parameter derivatives of the input Laplacian, see Equation (2). We have identified the direct children of W ( i ) in the Laplacian's compute graph, see Equation (D40). This allows us to compute the Jacobian J W ( i ) ∆ u by the chain rule, i.e. by accumulating the Jacobians over all direct children,

The terms ∇ · ∆ u can 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 to summation over the data) for W ( i ) as

The Gramian consists of nine different terms, see Figure D19 for a visualization which shows not only the diagonal blocks G ( i ) Ω , but also the full Gramian G Ω which decomposes in the same way. The terms ∇ · ∆ u are automatically computed when computing the gradient of the loss via backpropagation. We will now proceed and simplify the terms by inserting the Jacobians into Equation (D41) and studying the Gramian's block diagonal, which is approximated by KFAC, in more detail.

Computing J W ( i ) · Let us first compute the Jacobians J W ( i ) · in Equation (D41). The Jacobian of the linear layer's forward pass is

The Jacobian from the gradient backpropagation is

and the Jacobian from the Hessian backpropagation is

where K ∈ R dim( Z ) × dim( Z ) (denoting Z := W ⊤ XW ) is a permutation matrix that, when multiplied onto a vector whose basis corresponds to that of the flattened output Z , modifies the order from first-varies-fastest to last-varies-fastest, i.e.

Full interior Gramian

(forward, forward)

(forward, gradient)

(forward, Hessian)

(gradient, forward)

(gradient, gradient)

(gradient, Hessian)

(Hessian, forward)

(Hessian, gradient)

(Hessian, Hessian)

le-5

0

4

6

Re-introducing the layer indices, the expressions in Equation (D40) become

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 of 16 terms, as there are 4 summands from the Jacobians in Equation (D43). However, we can simplify into 9 terms:

First, ∇ 2 z ( i ) u is symmetric, that is

Second, we multiply the transpose Jacobian onto ∇ ∇ 2 z ( i -1) u ∆ u , which inherits symmetry from the Hessian, [ ∇ ∇ 2 z ( i -1) u ∆ u ] j,k = [ ∇ ∇ 2 z ( i -1) u ∆ u ] k,j . Due to this symmetry, the action of K (or K ⊤ ) does not alter it,

In other words, it does not matter how we flatten (first- or last-varies-fastest). This simplifies the VJP (last line in Equation (D42)) to

We can now write down the simplified Jacobian from Equation (D41), whose self-outer product forms the Gramian block for a linear layer's weight matrix,

where (1) is the contribution from the forward pass, (2) is the contribution from the gradient backpropagation, and (3) is the contribution from the Hessian backpropagation. The Jacobians from Equation (D43) allow to express the Gramian in terms of Kronecker-structured 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 according to Equation (D45) is that if we write out the Gramian's block G ( i ) Ω = J W ( i ) ∆ u (J W ( i ) ∆ u ) ⊤ , we obtain 9 terms of different structure. Defining a single Kronecker product approximation would involve introducing new approximations on top of those employed by Eschenhagen et al. [17]. Therefore, the forward Laplacian, or Taylor-mode, perspective we choose in the main text is advantageous as it allows to define KFAC without introducing new approximations.

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 of the paper that present them.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See the paragraph dedicated to limitations in §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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: Whereas we provide a derivation of the proposed Kronecker-factored approximation, 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 crossreferenced.

All assumptions should be clearly stated or 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 a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide a detailed description of the general experimental protocol in §A.1 and details for each individual experiment in §A, including all hyper-parameter search spaces and hyper-parameters 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 to 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 is a specific model and empirical evaluation, it 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 is 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to 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 ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: While we do not show mean and standard deviations over different model initializations for training curves, we aimed to use a thorough tuning protocol to avoid artifacts from insufficient hyper-parameter tuning and also conducted experiments with alternative hyper-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 one experiment in §A.9.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or the standard error of the mean.

It 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 a 96% CI, if the hypothesis of Normality of errors 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 of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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]

Justification: It is the goal of this paper to improve the accuracy of neural network-based PDE solvers. As the numerical solution of PDEs is a classic topic in 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe 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)?

Answer: [NA]

Justification: The paper does not release any data or models that have a high risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use 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 a 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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 are 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 not release new assets.

Researchers should communicate the details of the dataset/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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in 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 nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing 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 we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\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- 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?\n\nLIMITATIONS:\nThe authors addressed the limitations, but I think the answer in (7) Experiment Statistical Significance should be No, as there are no error bars.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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)\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nPlease see \"Weaknesses\"", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nEven 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.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}