Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks

Authors: felix  dangel
johannes  müller
marius  zeinhofer
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 11:24:24

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.

Main paper: 
 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] . 
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 = ∂bi /∂aj. 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 networku θ = 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.Let us consider a domain Ω ⊆ R d and the partial differential equation Lu = f in Ω , u = g on ∂Ω , with right-hand side f , boundary data g and a differential operator L, e.g. the negative Laplacian-Lu = ∆ x u = d i=1 ∂ 2 xi u.We parametrize u with a neural net and train its parameters θ to minimize the lossEQUATIONwith 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 GramianEQUATION(2) is the Gauss-Newton matrix of the residual r(θ) = ( rΩ(θ) ⊤ / √ NΩ, r ∂Ω (θ) ⊤ / √ N ∂Ω ) ⊤ ∈ R NΩ+N ∂Ω with interior and boundary residuals r Ω,n (θ) = Lu θ (x n ) -f (x n ) and r ∂Ω,n (θ) = u θ (x b n ) -g(x b n ).We review Kronecker-factored approximate curvature (KFAC) which was introduced by Heskes [22] , Martens & Grosse [38] in the context of maximum likelihood estimation to approximate the per-layer Fisher information matrix by a Kronecker product to speed up approximate natural gradient descent [1] . The Fisher associated with the loss 1 /2N N n=1 ∥u θ (x n )y n ∥ 2 2 with targets y n ∈ R isEQUATIONwhere 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. 2 The second regression loss can also include other constraints like measurement data.3 KFAC tackles the Fisher's per-layer block diagonal, F (θ) ≈ diag(F (1) (θ), . . . , F (L) (θ)) withF (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,F (l) (θ) ≈ 1 N N n=1 z (l-1) n z (l-1)n ⊤ =:A (l) ∈R h (l-1) ×h (l-1)EQUATIONIt 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• (θ), . . . , 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 Lu 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] . 
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 1) , applying the chain rule yields the propagation rule∂ xi u, ∂ 2 xi,xj u for i, j = 1, . . . , d. At the first layer, set z (0) = x ∈ R d , ∂ xi z (0) = e i ∈ R d , i.e., the i-th basis vector and ∂ 2 xi,xj z (0) = 0 ∈ R d . For a linear layer f θ (l) (z (l-1) ) = W (l) z (l-EQUATIONEQUATIONEQUATIONThe propagation rule through a nonlinear element-wise activation layer z (l-1) → σ(z (l-1) ) isEQUATIONEQUATION∂ 2 xi,xj z (l) = ∂ xi z (l-1) ⊙ σ ′′ (z (l-1) ) ⊙ ∂ xj z (l-1) + σ ′ (z (l-1) ) ⊙ ∂ 2 xi,xj z (l-1) ∈ R h (l) . (6c)Forward Laplacian For differential operators of special structure, we can fuse the Taylor-mode forward propagation of individual directional derivatives in Equations ( 5) and ( 6) and obtain a more efficient computation. E.g., to compute not the full Hessian but only the Laplacian, we can simplify the forward pass, which yields the forward Laplacian framework of Li et al. [29] . To the best of our knowledge, this connection has not been pointed out in the literature. Concretely, by summing (5c) and (6c) over i = j, we obtain the Laplacian forward pass for linear and activation layersEQUATIONEQUATIONThis 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 xi,xj , the forward pass for a linear layer is Lz (l) = W (l) Lz (l-1) ∈ R h (l) , generalizing (7a), and similarly for Equation (7b)Lz (l) = σ ′ (z (l-1) ) ⊙ Lz (l-1) + d i,j=1 c i,j σ ′′ (z (l-1) ) ⊙ ∂ xi z (l-1) ⊙ ∂ xj z (l-1) ∈ R h (l) ,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.Let's consider the Poisson equation's interior Gramian block for a linear layer (suppressingΩ in N Ω ) G (l) Ω (θ) = 1 N N n=1 (J W (l) ∆ x u n ) ⊤ J W (l) ∆ x u n .Because we made the Laplacian computation explicit through Taylor-mode autodiff ( §3.1, specifically Equation (7a)), we can stack all output vectors that share the layer's weight into a matrixZ (l) n ∈ R h (l) ×S with S = d + 2 and columns Z (l) n,1 = z (l) n , Z (l) n,2 = ∂ x1 z (l) n , . . . , Z (l) n,1+d = ∂ x d z (l)n , andZ (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 ruleJ W (l) ∆ x u n = (J Z (l) n ∆ x u n )J W (l) Z (l) n = S s=1 Z (l-1) n,s ∈R h (l-1) ⊤ ⊗ J Z (l) n,s ∆ x u n =:g (l) n,s ∈R h (l), which has a structure similar to the Jacobian in §2.2, but with an additional sum over the S shared vectors. With that, we can now express the exact interior Gramian for a layer asEQUATIONNext, 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 operatorEQUATIONTo generalize the previous section, let's consider the general M -dimensional PDE system of order k,EQUATIONwhere D m x u collects all partial derivatives of order m. For m ∈ {0, . . . , k} there are S m = d+m-1 d-1independent partial derivatives and the total number of independent derivatives isS := 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 (θEQUATIONThe 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 formEQUATIONwith Λ(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 findEQUATIONfor any m. Hence, we can derive a forward propagation for the required derivatives where a linear layer processes at most S vectors3 , 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 1x z (l-1) , . . . , D k x z (l-1) ) ∈ R h (l-1) ×S (andZ (l)n for the outputs), the chain rule yieldsG (l) Ω (θ) = 1 N N n=1 J W (l) Z (l) n ⊤ J Z (l) n r Ω,n ⊤ Λ(r Ω,n ) J Z (l) n r Ω,n J W (l) Z (l) n = 1 N N,S,S n,s,s ′ =1 J W (l) Z (l) n,s ⊤ J Z (l) n,s r Ω,n ⊤ Λ(r Ω,n ) J Z (l) n,s ′ r Ω,n J W (l) Z (l) n,s ′ = 1 N N,S,S n,s,s ′ =1 Z (l-1) n,s Z (l-1)⊤ n,s ′ ⊗ J Z (l) n,s r Ω,n ⊤ Λ(r Ω,n ) J Z (l) n,s ′ r Ω,nwhereZ (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 operatorEQUATIONTo 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 formEQUATIONwhere ℓ : 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] .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)), GEQUATIONExponential moving average and damping For preconditioning, we accumulate the Kronecker factors A (l)EQUATION•,t over time using an exponential moving average Â(l)•,t = β Â(l) •,t-1 + (1 -β)A (l) •,tof 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, Ã(l)•,t = Â(l) •,t +λI and B(l) •,t = B(l)•,t +λI such that the curvature approximation used for preconditioning at step t isG •,t ≈ diag Ã(1) Ω,t ⊗ B(1) Ω,t , . . . , Ã(L) Ω,t ⊗ B(L) Ω,t + diag Ã(1) ∂Ω,t ⊗ B(1) ∂Ω,t , . . . , Ã(L) ∂Ω,t ⊗ B(L) ∂Ω,t .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 = -( Ã(l) Ω,t ⊗ B(l) Ω,t + Ã(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 [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] .α ⋆ µ ⋆ = - ∆ ⊤ t G(θ t )∆ t + λ ∥∆ t ∥ 2 ∆ ⊤ t G(θ t )δ t + λ∆ ⊤ t δ t ∆ ⊤ t G(θ t )δ t + λ∆ ⊤ t δ t δ ⊤ t G(θ t )∆ t + λ ∥δ t ∥ 2 -1 ∆ ⊤ t g t δ ⊤ t g t (seeComputational 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.  
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] [41] to reproduce ENGD's performance (Figure 1 ). It is given byEQUATION)We choose a fixed data set of same size as the original paper, then use random/grid search to evaluate the performance of all optimizers for different tanh-activated MLPs, one shallow and two with five fully-connected layers of different width (all details in §A.2). We include ENGD whenever the network's parameter space is small enough to build up the Gramian.For the shallow net (Figure 1 , left), we can reproduce the results of [41] , where exact ENGD achieves high accuracy. In terms of computation time, our KFACs are competitive with full-ENGD for a long phase, outperforming the first-order and quasi-Newton baselines. In contrast to ENGD, which runs out of memory for networks with more than 10 000 parameters, KFAC scales to larger networks (Figure 1 , center and right) and is competitive with other second-order optimizers like Hessian-free, which uses more sophisticated heuristics. We make similar observations on a small (1+1)d heat equation with the same models, see §A.7 and fig. A10 .To demonstrate that our methods can also be applied to other problems than the Poisson equation, we consider a four-dimensional heat equationEQUATIONwith 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 = 116 864 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 .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 solutionsEQUATIONWe 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 EQUATION 
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] . D 
 
 Appendix: 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.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 conditionsu(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 -Hessian-free: curvature matrix: GGN, initial damping: 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): 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 Ω = 3 000, 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 -LBFGS: learning rate: • Hessian-free: curvature matrix: U({GGN, Hessian}), initial damping: . We assign each run a budget of 10 000 s.• 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 ]) -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}) -LBFGS: learning rate: U({5 • 10 -1 , 2 • 10 -1 , 1 • 10 -1 , 5 • 10 -2 , 2 • 10 -2 , 1 • 10 -2 }),-SGD: learning rate: LU([1 • 10 -3 ; 1 • 10 -2 ]), momentum: U({0, 3 • 10 -1 , 6 • 10 -1 , 9 • 10 -1 }) -Adam: learning rate: LU([1 • 10 -4 ; 5 • 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 }),-SGD: learning rate: LU([1 • 10 -3 ; 1 • 10 -2 ]), momentum: U({0, 3 • 10 -1 , 6 • 10 -1 , 9 • 10 -1 }) -Adam: learning rate: LU([1 • 10 -4 ; 5 • 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 }),U({5 • 10 -1 , 2 • 10 -1 , 1 • 10 -1 , 5 • 10 -2 , 2 • 10 -2 , 1 • 10 -2 }),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 }),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 -SGD: learning rate: 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):• ∂ t u(t, x) -κ∆ xu(t, x) = 0with 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∂ x1 u(x) -κ d d=2 ∆ x d u(x) = 0 .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 isu ⋆ (t, x) = exp - π 2 dt 4 d i=1 sin(π[x i ]) .The PINN loss for this problem consists of three terms: a PDE term, an initial value condition term, and a spatial boundary condition term,L(θ) = 1 N Ω NΩ n=1 ∂ t u θ (x Ω n ) - 1 4 ∆ xn u θ (x Ω n ) 2 + 1 N ∂Ω N ∂Ω n=1 u θ (x ∂Ω n ) -g(x ∂Ω n ) 2 + 1 N 0 N0 n=1 u θ (0, x 0 n ) -u 0 (x 0 n ) 2 withx Ω 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 = Ncond /2 without loss of generality, we writeL(θ) = 1 N Ω NΩ n=1 ∂ t u θ (x Ω n ) - 1 4 ∆ xn u θ (x Ω n ) -y Ω n 2 2 LΩ(θ) + 1 N cond Ncond n=1 u θ (x cond n ) -y cond n 2 2Lcond(θ)with domain inputs x Ω n ∼ T × Ω and targets y Ω n = 0, boundary and initial condition targetsy cond n = u ⋆ (x cond n ) with initial inputs x cond n ∼ {0}× Ω for n = 1, . . . , Ncond /2 and boundary inputs x cond n ∼ T × ∂Ω for n = Ncond /2 + 1, . . . , N cond . This loss has the same structure as the PINN loss in Equation (1).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 -LBFGS: learning rate: U({5 • 10 -1 , 2 • 10 -1 , 1 • 10 -1 , 5 • 10 -2 , 2 • 10 -2 , 1 • 10 -2 }),-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 ]) -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}) -LBFGS: learning rate: U({5 • 10 -1 , 2 • 10 -1 , 1 • 10 -1 , 5 • 10 -2 , 2 • 10 -2 , 1 • 10 -2 }),({5 • 10 -1 , 2 • 10 -1 , 1 • 10 -1 , 5 • 10 -2 , 2 • 10 -2 , 1 • 10 -2 }),4 i=1 sin(2x i ) for t, x ∈ [0, 1] × ∂[0, 1]4 throughout time, and initial value conditions u(0, x) = 4 i=1 sin(2x i ) for x ∈ [0, 1] 4 . We sample training batches of size N Ω = 3 000, 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(2x 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 = 10 065 trainable weights, and a five layer 5 → 256 → 256 → 128 → 128 → 1 MLP with D = 116 864 trainable weights. For the biggest architecture, full and layer-wise ENGD lead to out-of-memory errors and are thus not tested. Figure A12 visualizes the results. -LBFGS: learning rate: 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.U({5 • 10 -1 , 2 • 10 -1 , 1 • 10 -1 , 5 • 10 -2 , 2 • 10 -2 , 1 • 10 -2 }),Best run details The runs shown in Figure A14 correspond to the following hyper-parameters:• SGD: learning rate: • Hessian-free: curvature matrix: GGN, initial damping: 5.077 634 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 ∂ t p + div(µp) - 1 2 Tr(σσ ⊤ ∇ 2 p) = 0, p(0) = p 0 , 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∂ t q + div(µ) + ∇q • µ - 1 2 ∥σ ⊤ ∇q∥ 2 - 1 2 tr(σσ ⊤ ∇ 2 q) = 0, q(0) = log p 0 .For the concrete example of the main text, we set µ(t, x) = -1 2 x and σ = √ 2I ∈ 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 ∂ t q(t, x) - d 2 - 1 2 ∇q(t, x) • x -∥∇q(t, x)∥ 2 -∆q(t, x) = 0, q(0) = log(p * (0)),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.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 pointx ∈ R d along a direction αv ∈ R d with α ∈ R, f (α) = f (x + αv) = f (x) + α ∂f (x) ∂x ⊤ v + α 2 2! v ⊤ ∂ 2 f (x) ∂x 2 v + α 3 3! i1,i2i3 ∂ 3 f (x) ∂x 3 i1,i2,i3 v i1 v i2 v i3 + . . . + α K K! i1,...,i K ∂ K f (x) ∂x K i1,...,i K v i1 • • • v i K .We can unify this expression by introducing the K-th order directional derivative of f at x along v,∂ K f (x) [v, . . . , v] K times := i1,...,i K ∂ K f (x) ∂x K i1,...,i K v i1 . . . v i K .This simplifies the uni-directional Taylor expansion tof (α) = f (x + αv) = f (x) + α∂f (x)[v] + α 2 2! ∂ 2 f (x)[v, v] + α 3 3! ∂ 3 f (x)[v, v, v] + . . . + α K K! ∂ K f (x)[v, . . . , v] =: K k=1 α k k! ∂ k f (x) ⊗ k v =: K k=1 w f k α kwhere 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 wf k ∈ R c such that [w f k ] i is the k-th order Taylor coefficient of the ith component of f .A note on generality: In this introduction to Taylor-mode, we limit the discussion to the 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. inside the forward propagation scheme, specifically Equation (C24c), and push-forward the summed second-order coefficients. This simplifies Equation (C24) toEQUATIONEQUATIONEQUATIONEquation (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.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 asJ 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:EQUATIONEQUATIONEQUATIONNow, let g : R Din → R Dout be a linear layer with weight matrix W ∈ R Dout×Din and bias vector b ∈ R Dout . Its Jacobian is J h(x) (W h(x) + b) = W and the second-order derivative is zero. Hence, Equation (C25) for linear layers becomesEQUATIONEQUATIONEQUATIONWe 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:w f 0 w f 1,1 . . . w f 1,d D i=1 w f 2,i :=T f ∈R D out ×(d+2) = (W b) w h 0 w h 1,1 . . . w h 1,d d i=1 w h 2,i 1 0 . . . 0 0 :=T h ∈R (D in +1)×(d+2), or, in compact form,T f = W T h . (C28)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 Din×(d+2) during the computation of the net's prediction and Laplacian via the forward Laplacian framework and yields another matrix T (l) ∈ R Dout×(d+2) .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)EQUATIONHere, part({1, . . . , K}) denotes the set of all set partitions of {1, . . . , K} (σ is a set of sets). E.g.,part({1}) = {{{1}}} , part({1, 2}) = {{{1, 2}}, {{1}, {2}}} , part({1, 2, 3}) = {{{1, 2, 3}}, {{1}, {2, 3}}, {{1, 2}, {3}}, {{1, 3}, {2}}, {{1}, {2}, {3}}} .To make this more concrete, let's consider Equation (C29) for first-and second-order derivatives,EQUATIONEQUATIONFrom Equation (C30), we can see that if we want to compute a weighted sum of second-order derivativesi,j c i,j ∂ 2 f (x)[v i , v j ],we can pull the sum inside the second equation,EQUATIONHence, 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 thusEQUATIONEQUATIONEQUATIONThe 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 tensorEQUATIONEQUATION) 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 layersEQUATIONEQUATIONand for activation layersEQUATIONEQUATION(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.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:1. Autodiff Laplacian. Computes the Laplacian with PyTorch's automatic differentiation (functorch) by computing the batched Hessian trace (via torch.func.hessian and torch.func.vmap). This is the standard approach in many PINN implementations. 2. Forward Laplacian. Computes the Laplacian via the forward Laplacian framework. We used this approach for all PDEs and optimizers, except ENGD, presented in the experiments.We use the biggest network from our experiments (the 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).The following tables compare run time and peak memory between the two approaches: 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 Ω Autodiff ForwardHere, 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 ∇ 2x 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.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 ∇ z (i-1) u = J z (i-1) z (i) ⊤ ∇ z (i) u , Hessian backpropagation Just like gradient backpropagation, we can derive a recursive scheme for the Hessian. Recall the Hessian chain ruleEQUATIONEQUATIONwhere 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)∇ 2 z (i-1) u = J z (i-1) z (i) ⊤ ∇ 2 z (i) u J z (i-1) z (i) + h (i) k=1 ∇ 2 z (i-1) [z (i) ] k [∇ z (i) u] k , ∇ 2 θ (i) u = J θ (i) z (i) ⊤ ∇ 2 z (i) u J θ (i) z (i) + h (i) k=1 ∇ 2 θ (i) [z (i) ] k [∇ z (i) u] k (D37)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 ∇ 2x u = ∇ 2 z (0) u: Re-introducing the layer indices, the expressions in Equation (D40) becomeEQUATION)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 isJ W (i) W (i) ⊤ ∇ 2 z (i) u W (i) = I ⊗ W (i) ⊤ ∇ 2 z (i) u + K I ⊗ W (i) ⊤ ∇ 2 z (i) u ,and the transposed Jacobian isI ⊗ ∇ 2 z (i) u W (i) + I ⊗ ∇ 2 z (i) u W (i) K ⊤ .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,K ⊤ ∇ ∇ 2z (i-1) u ∆u = ∇ ∇ 2 z (i-1) u ∆u . In other words, it does not matter how we flatten (first-or last-varies-fastest). This simplifies the VJP (last line in Equation (D42)) toI ⊗ ∇ 2 z (i) u W (i) ∇ ∇ 2 z (i-1) u ∆u + I ⊗ ∇ 2 z (i) u W (i) K ⊤ ∇ ∇ 2 z (i-1) u ∆u = 2 I ⊗ ∇ 2 z (i) u W (i) ∇ ∇ 2z (i-1) u ∆u . 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,EQUATIONEQUATION+ 2 I ⊗ ∇ 2 z (i) u W (i) ∇ ∇ 2 z (i-1) u ∆u,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.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.• 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. • 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.Depending on the linear operator, one may reduce weight sharing, as demonstrated for the Laplacian in §3.1. 


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

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

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

Reviews: review: paper_summary: paper_summary: This work considers the problem of optimising partial differential equations (PDE) with neural networks, in particular second-order optimization.
Even 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.
strength_weakness: strength_weakness: Even if experiments with a relatively large number of parameters are shown (10^5), the paper does not provide the complexity of the algorithm, in terms of relevant quantities, for example the input dimension, output dimension, batch size, number of parameters, etc. For example, does the method scale poorly with the input dimension? I would expect quadratic scaling with input dimension. I have not read the forward laplacian paper, but I have the feeling that improves scaling by only a constant factor. https://arxiv.org/abs/1206.6464 argued that it may not exist an algorithm for computing the diagonal of the Hessian that is linear in the dimension using automatic differentiation. Also, as with other KFAC approximations, it remains unclear how general is the approach, e.g. how to apply it to higher order PDEs or other neural network models.
questions: questions: NA
limitations: limitations: NA
paper_summary: paper_summary: This paper generalizes the K-FAC method (which is well-known in optimization for deep learning) to enable it to train PINNs. The key idea is to combine Taylor-mode autodifferentiation with recent work on K-FAC with weight sharing layers. The proposed method is evaluated on several PDEs, where it outperforms first-order methods like SGD and Adam and quasi-Newton methods like L-BFGS, while performing comparably to (matrix-free) ENGD.
strength_weakness: strength_weakness: * The paper is missing more recent work on difficulties of training PINNs, such as [1, 2].

* K-FAC/K-FAC* do not reach the lowest error obtained by running ENGD with $D = 449$ (Figure 2). Why is this the case?

* 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]?

* 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?

* 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)})$.

* I’d recommend defining the symbol $\odot$ as the Hadamard product before it appears in equation (6).

* Equation below line 168: $N_\Omega$ in the sum should be $N$.

* I believe the outer product decomposition in line 200 is missing a transpose on the second $l_{n, m}$ in the summation.

* There is no per-step computational complexity provided for the proposed methods.

* No open-source implementation is provided for the proposed methods.

Minor comments:

* 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). 

* For consistency with existing literature, I would recommend writing “KFAC” as “K-FAC” and “LBFGS” as “L-BFGS”.

[1] “On the Role of Fixed Points of Dynamical Systems in Training Physics-Informed Neural Networks.” Rohrhofer et al. (2023)

[2] “Challenges in Training PINNs: A Loss Landscape Perspective.” Rathore et al. (2024)

[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)

[4] “When and why PINNs fail to train: A neural tangent kernel perspective.” Wang et al. (2022)

[5] “PINNacle: A Comprehensive Benchmark of Physics-Informed Neural Networks for Solving PDEs.” Hao et al. (2023)

[6] “PDEBENCH: An Extensive Benchmark for Scientific Machine Learning.” Takamoto et al. (2022)
questions: questions: * Line 64: Could the authors please elaborate on why the preconditioner not having a Laplacian term is beneficial?

* 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?

* 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.

* 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?

* Are there operators besides the Laplacian for which the weight sharing can be reduced?
limitations: limitations: Please see "Weaknesses"
paper_summary: paper_summary: This work develops a KFAC approximation for approximate second order optimization of physics-informed neural network (PINN) losses, in order to address optimization difficulties of PINNs.The first key idea is to use forward mode automatic differentiation for the different derivatives of the PDE and notice that it corresponds to a forward pass with additional inputs (corresponding to the different derivatives) and weight sharing. The second key idea is to then tie this to KFAC approximations for weight sharing architectures. 
Both 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.
strength_weakness: strength_weakness: - 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.). 

- 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. 

---------------------
- 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.
questions: questions: - I noticed that the EMA update rate hyper-parameter is often super large, even up to 0.988 (for heat equation, appendix A9), which effectively looks only at the most recent batch. Do you have an explanation for why this is? What happens if you use a much smaller mini-batch size?
- 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? 
- 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? 
- It seems very odd that KFAC performs better if the random mini-batch are *not* sampled (new) at every step. Have you noticed any such thing for KFAC on standard problems (not PINN losses)? Are the matrices in KFAC inverted every iteration or also just every T-th iteration?
limitations: limitations: The authors addressed the limitations, but I think the answer in (7) Experiment Statistical Significance should be No, as there are no error bars.

Comments: title: Thanks for the rebuttal
Answer to Rebuttal
Additional experiment with error bars
comment: Thank you for your question. Assume we are given an inverse problem (let's say a source recovery problem for simplicity) where we are looking to find both the function $u$ and the right-hand side/source $f$ given some observations of $u$ and a PDE that $u$ should satisfy. In this case, the loss function contains a term for the PDE residual and boundary values, and additionally a data fidelity term that incorporates the observations of $u$. A canonical approach would be to parametrize both $u$ and $f$ by neural networks, which explains the two neural networks in our answer above.
Thank you for addressing most of my comments. I would have liked to see the results on the log-Fokker-Planck equation, but it is ok if this is not available in time (I understand the discussion period is rather short). I will raise my score accordingly.

Minor question (does not affect my score): why would solving an inverse problem with this approach require two (or more) neural networks? I thought the standard formulation for the PINN inverse problem uses a single neural network with the inverse coefficients as inputs.
Thank you for your response and for answering my questions.

> We respectfully disagree. We report relative L2 error rather than training loss [...]

Thank you for the clarification, this was a misunderstanding on my side. This is indeed mentioned in Sec. 4. 

> It is true that our plots show the performance of a single run. However, this run was obtained through extensive hyper-parameter tuning [...]

I was about to suggest running 10 different seeds and data splits, which is what you have reported in the comment below. 

> High EMA factors correspond to slowly forgetting pre-conditioners [...]

Thanks, this makes sense, of course, it was a misunderstanding on my side.

> We found that for all optimizers it was favorable to keep a mini-batch for significant number of iterations.

This is quite interesting and weird at the same time. If you find the time, an ablation for this would be quite interesting.
Dear Reviewer aiiA,

We are happy to inform you that we can report the results of an additional experiment testing the optimizers for a variety of initializations. For this, we took the optimized hyper-parameters and run the different optimizers for 10 different initializations and mini-batches for the 4+1d heat equation for the MLP with 10 065 parameters (Figure 2, middle plot). Unfortunately, we are not allowed to share figures during the discussion period. The table below summarizes variability of final performance. We will add the corresponding figure to the appendix. 

We find that all optimizers except LBFGS perform stably when using different data/model initialization. When taking the average performance, the same trend as shown in the original submission is visible where KFAC* outperforms the competing optimizers. Please let us know in case that you have any further questions. 


| Optimizer | Relative L2 error | 
|-----------|--------------------|
| SGD |  $(6.77\pm2.36)\cdot 10^{-3}$ | 
| Adam | $(3.10\pm2.45)\cdot 10^{-3}$ | 
| Hessian-free | $(1.97\pm0.47)\cdot 10^{-5}$ |
| ENGD (full) | $(4.99\pm9.25)\cdot 10^{-2}$ | 
| ENGD (layer-wise) | $(1.68\pm0.26)\cdot 10^{-4}$ |  
| KFAC | $(6.50\pm0.84)\cdot 10^{-5}$ | 
| KFAC* | $(1.46\pm0.48)\cdot 10^{-5}$ |

Let us know if you would like to see similar results for other sub-experiments.
Thanks for your follow-up question!

You are of course right that the MLP's cost depends linearly on $d_\Omega$ due to the first layer's matrix multiply ($d_\Omega \to 768$). For the values of $d_\Omega$ we experimented with, the evaluation cost is dominated by the sub-sequent matrix multiplies (e.g. $d_\Omega \le 100$ but the second layer is $768 \to 768$). Therefore we empirically observe linear scaling in this regime, i.e. 'approximately' no dependency of the MLP w.r.t. $d_\Omega$ and linear scaling from the number of forward passes for the Laplacian.

You are completely right though that if we let $d_\Omega \to \infty$, e.g. for score matching with $d_\Omega \sim 10^6$, the empirical behavior would eventually become quadratic on this model. However, for the PDE applications we target in this paper, we think that practically relevant values are $d_\Omega \sim 10^2$ (at most $10^3$). Therefore, we believe the empirical linear scaling we measured in this regime is representative for our method's scaling on PINN problems.

We will add a sentence to clarify this limitation for extremely high dimensions of $d_\Omega$ which are (1) out of scope for PINNs and (2) for which the autodiff-based approach suffers from the same limitation.

Please let us know if you have any follow-up question.
Thank you for the response.
How is the forward pass independent on the input dimension?
In an MLP, you must compute a matrix-vector product, where the input vector has size d, and the matrix has size d' x d (hidden layer size d'), this has complexity o(d'd), linear in d. Maybe d is not your bottleneck now, but if you increase it at some point you should suffer from it.
I'm asking also because I have been working on score matching, which requires computing the trace of the Hessian of the log-likelohood, different setting but same computational problem (see for example https://arxiv.org/abs/1905.07088). That is impossible to scale up when d is in the range of millions or higher.