Improving Energy Natural Gradient Descent through Woodbury, Momentum, and Randomization
Authors: Andres Guzman-Cordero
Felix Dangel
Gil Goldshlager
Marius Zeinhofer
Venue: NeurIPS 2025 poster ()
Keywords: Physics-Informed Neural Networks
Natural Gradient Descent
Woodbury Matrix Identity
Momentum
Nyström Approximation
Primary area: machine_learning_for_sciences
TL;DR: We introduce Woodbury's matrix identity, momentum-like SPRING and randomization to make energy natural gradient descent 75 times faster for PINNs.

Abstract: Natural gradient methods significantly accelerate the training of Physics-Informed Neural Networks (PINNs), but are often prohibitively costly. We introduce a suite of techniques to improve the accuracy and efficiency of energy natural gradient descent (ENGD) for PINNs. First, we leverage the Woodbury formula to dramatically reduce the computational complexity of ENGD. Second, we adapt the Subsampled Projected-Increment Natural Gradient Descent algorithm from the variational Monte Carlo literature to accelerate the convergence. Third, we explore the use of randomized algorithms to further reduce the computational cost in the case of large batch sizes. We find that randomization accelerates progress in the early stages of training for low-dimensional problems, and we identify key barriers to attaining acceleration in other scenarios. Our numerical experiments demonstrate that our methods outperform previous approaches, achieving the same $L^2$ error as the original ENGD up to $75\times$ faster.

Improving Energy Natural Gradient Descent through Woodbury, Momentum, and Randomization
Andrés Guzmán-Cordero
Vector Institute, Mila Quebec AI Institute, Université de Montréal
andres. wndres.gumaan-corooen
quebec
Felix Dangel
Gil Goldshlager
Vector Institute
UC Berkeley
f. w.mngeeevectorrnsstttueeee.
ggoldshh@errkeey.com edu
Marius Zeinhofer
ETH Zurich
marius. marius.zeinhofeee www.zzzeennnneeeeeeeo math. ethz. ch
Abstract
Natural gradient methods significantly accelerate the training of Physics-Informed Neural Networks (PINNs), but are often prohibitively costly. We introduce suite of techniques to improve the accuracy and efficiency of energy natural gradient descent (ENGD) for PINNs. First, we leverage the Woodbury formula dramat- ically reduce the computational complexity of ENGD. Second, we adapt the Sub- sampled Projected-Inereee Increment Natural Gradient Descent algorithm from the varia- tional Monte Carlo literature to accelerate the convergence. Third, we explore the use of randomized algorithms to further reduce the computational cost the case of large batch sizes. We find that randomization accelerates progress in the early stages of training for low-dimensional problems, and we identify key barriers attaining acceleration in other scenarios. Our numerical experiments demonstrate that our methods outperform previous approaches, achieving the same L error as the original ENGD up to 75x faster.
1 Introduction
Using neural networks to solve partial differential equations (PDEs) is a promising and highly active research direction at the intersection of machine learning and scientific computing. While the idea (PINNs)
is not new 7 20), has gained traction with the advent of Physics Informed Neural Networks
3 Compared traditional methods, neural networks avoid the need for sophisticated, problem-dependent discretization schemes and promise to scale better to high-dimensional problems 33 However, it is very challenging attain the accuracy levels of traditional solvers with neural networks. This is in large part due to the loss landscape's non-convexity and ill-conditioning [44), To address this problem, second-order methods such energy natural gradient descent [ENGD. 27 and Cronecker-factored approximate curvature [KFAC, 24 have been proposed. These methods greatly improve upon the accuracies attainable by first-order optimizers like SGD or Adam (6), but they suffer from either high per-iteration cost or high implementation complexity. As result, there remains significant need for more accurate and efficient methods to train PINNs.
In parallel to the PINN community, neural networks have also been leveraged in conjunction with the variational Monte Carlo method [VMC. 9 to simulate quantum many-body systems (3] [4] 21 32
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

SPRING 11
VMC+SR* 9
NNs in VMC
K-FAC 32
MinSR 4
K- F-C 6
Sketch and Precond.
PINNS 38
ENGD 27
Us
37
Nystrom Mothod
Sketch and Solve? (38)
Low-rank Approx. 133
Nyströmm PCG
48
2000
2001
2012
our
ah
caue
Cznz
PZEE
Figure I: Timeline of VMC, PINN and RLA methods. *: Stochastic reconfiguration (SR) and energy natural gradient descent (ENGD) both precondition their stochastic gradients with the inverse of an appropriate curvature matrix. f: Inspired by MinSR, we apply Woodbury's matrix identity to use the kernel matrix instead of ENGD's Gauss-Newton matrix, thus reducing the cost to O(N2 P) instead of O ((P)). where N denotes the batch size and P the number of parameters; we further introduce SPRING for PINNs to transport curvature information across optimization steps; last, we use GPU-efficient Nyström approximation to reduce the iteration cost for large batch sizes.
Second-order optimizers are very common within this field and the most popular one, known as stochastic reconfiguration [2], shares similar computational structure ENGD, owing to similar mathematical derivation as a projected functional algorithm 288.. Within the neural network VMC community, the quest model complex quantum systems to as high of accuracy as possible has led to significant advances in algorithmic implementations of the SR optimizer, including the use of the Woodbury matrix identity reduce the computational cost of each optimization step [MinSR, 4 and the development of momentum schemes that can accelerate the convergence of natural gradient methods [SPRING,
Given the similarities between VMC and PINNs is natural to wonder if PINNs can benefit from the techniques developed in the more mature optimization literature for VMC. In this paper, answer this question in the affirmative (See Figure D. We first show that both MinSR and SPRING can be adapted drastically improve ENGD for PINNs. Going beyond transfer, explore the of randomized methods to further reduce the cost of both ENGD and SPRING iterations. Concretely, our contributions are as follows:
We transfer ideas from the VMC community to PINNs, showing that the Woodbury matrix identity drastically improves the computational bottleneck of ENGD for PINNs. Concretely (details in Section 3). we compute in sample space, as its dimensionality is typ- ically much smaller than the parameter space. This enables exact ENGD for PINNs with up to 10" trainable weights on standard consumer hardware with drastically improved com- putation time. Unlike K-FAC, the implementation flexible. We then show that the SPRING algorithm introduced in the VMC community can further improve the training process of PINNs, yielding consistently faster convergence and removing the need for the expensive line search from ENGD. The resulting method, summarized in algorithm [ is straightforward to implement and achieves state-of-te sate-of-the-as -hthhe--r art results. and thus recommend it as the new optimizer of choice for training PINNs.
aechitecture-agnossee and thus highly
Going beyond transfer, we explore the use of randomized algorithms based on a variant Nyström's 's method [10 to further accelerate both ENGD and SPRING in the case of large sample sizes. As standard formulations of Nyström's 's method require SVDs, which are ex- tremely slow on GPU in contemporary ML libraries, we develop a GPU-efficient variant of Nyström that only uses Cholesky solves on square matrices of sketch dimension. We high- light this GPU-efficient Nyström approximation as an independent contribution which can be useful in a wide array of machine learning applications. For PINNs, we find that it can accelerate the early phases of training for low-dimensional problems, and we identify key barriers to extending this acceleration to the later phases and nigher-dimensione settings.

---

We present thorough empirical evidence for our claims, demonstrating the utility of the Woodbury identity and the SPRING algorithm as well the power and the limitations of randomization. We explore both low- and high-dimensional problems and we meticulously tune the hyperparameters of each method to aftain optimal performance, with final results supported by more than 4500 training runs in total.
Related work The Woodbury identity has also been used train neural networks for tradi- tional machine learning problems [5] 49 35 and other applications, such as model sparsification 4Л, involving the Fisher information matrix in traditional deep learning. Descriptions of the ill- conditioned landscape of PINNs and some solutions have been explored Rathore et al. (4 While the key contribution of Rathore et
4 addresses the scalability of Newton's 's method large neu- ral networks and is thus similar to our results, their focus is on Newton's s method rather than natural gradient methods. Kiyani et al. [88 evaluate quasi-Newton Newton methods which leverage historical gradi- ents for improved efficiency and accuracy across stiff and non-linear PDEs. Wang et al. [46) provides a theoretical framework for gradient alignment in multiobjective PINN training, demonstrating the potential of second-order methods to resolve directional conflicts through Hessian preconditioning.
2 Background
Physics-Informet Neural Networks PINNs reformulate forward or inverse PDE problems as least-squares minimization problems. Constraints, such as boundary conditions data terms are typically included into the loss function as soft penalties. We illustrate the approach using a general PDE operator L and Dirichlet boundary conditions. Suppose aim to solve the PDE
Lu in Ω on 82.
Here Z C. IRd is the computational domain (possibly including time). f is a forcing term and g are the boundary values. Introducing neural network ansatz uo with trainable parameters R the above equation is reformulated as least-squares minimization problem
|0
No
[82]
Nao
L( =
Σ (Cuo((x; - f(x,))² + 2N
2N Σ
(ua(x))-9aa))
1=1
where Ti,.... No are randomly drawn points in Ω and .f...................... Noo are randomly drawn points on OΩ. Note that it holds L(f) = 0, then the PDE is satisfied at the points in the interior and the boundary conditions met the boundary points.
Natural Gradient Methods for PINNS It is well-known that PINNs are notoriously difficult to train with first-order methods 44 45) Natural gradient methods present a promising alternative, able to achieve excellent accuracy far beyond what is possible with first-order methods or quasi- Newton methods (28 27 16 6). The natural gradient descent scheme introduced in 277, which we refer to as energy natural gradient descent is of the form
θ 02 - nr (Ox) + LL))), = 0,1,,
(1)
where 7/s is well chosen step-size and À > is regularization parameter. The matrix G(Ok) is the Riemannian metric induced by the functional Hessian in the tangent space of the ansatz, given by
[221
No Jø (ug Jo Lug (zx)
[anl
Nan
G(0)= = No
+
Σ Ja us (x("))) Ja us (x))),
(2)
Non
n=1
where Jo denotes the Jacobian with respect to the trainable parameters. Here, we assumed the PDE operator C to be linear, for a nonlinear PDE operator we use its linearization, compare also to [16].
Scalability of Natural Gradient Methods for PINNs The matrix G(Ox) E PxP is quadratic in number trainable network parameters. Hence, naive approaches to solve G(O,)xr = (L(O2) are doomed to fail for larger networks due to the cubic cost O P3).. Strategies to treat the system in regimes where direct solvers are infeasible include matrix-free approaches based on the conju- gate gradient method [16] 23 with matrix-vecte products 31 40), and approximating G(Og) with
3

---

an casy-to-invert strategy suffers from the ill-conditioning of the matrix G(O+) 44 and can require large amounts of CG iterations to produce an acceptable solution to the linear system. The K-FAC approach [6] is complex and depends on the network architecture and the PDE. As described in detail Section we exploit different structure in G' using low-rank factorization present in the matrix G(Ox). ) and
[6] Both approaches have drawbacks. The matrix- free
we show that it yields a fast and accurate method to solve G(Og)a (r = ((((s) that can naturally be combined with powerful techniques of randomized linear algebra.
To improve the strategies of matrix-free solvers and K-FAC, we draw inspiration from the field of neural network wavefunctions, which leverages the variational Monte Carlo method [9] to simulate quantum systems from lattice models [3] to molecules and materials Almost all works on neural network wavefunctions are based on an optimization approach known as stochastic reconfiguration (SR) 42), which is quantum generalization of natural gradient descent (see [32], Appendix C). In fact, the same eq. [ applies for SR, but in this case G(& the Fisher information matrix. A number of works have demonstrated that such approaches outperform simpler first-order optimization strategies when training neural network wavefunctions [32 22 3g
Variational Monte Carlo
14
21.
Recently, Chen & Heyl 4 proposed the MinSR method to tackle the scalability of SR in the VMC context. The MinSR algorithm is able to directly calculate the SR search direction with a drastically reduced computational cost, relying the fact that the rank of ((Ok) is upper bounded by the batch
size N. Even more recently, Goldshlager et al. [T] proposed the SPRING algorithm to accelerate the convergence of MinSR by incorporating special form of momentum. In Section 3 how these methods can be applied in the context of PINNs.
we show
VMC meets PINNs. It is worth noting that variational Monte Carlo applications, samples are drawn from a special probability density 24 p(x) where p(x:) depends on the wavefunction itself. However, this does not change the fundamental structure of the algorithms or the transferability of techniques from one domain to the other. The unifying principle is that in both application domains a single probability distribution is used to define the loss function, the natural gradient direction, and the sampling algorithm. The only difference is that in variational Monte Carlo the distribution happens to be dependent on the wavefunction, whereas for PINNs the distribution is a weighted mixture of the uniform distributions over the interior and the boundary.
Sketching Randomized algorithms have made major impact on numerical linear algebra over the last two decades, providing faster algorithms for many important problems including low-rank matrix approximations, linear systems, and eigenvalue solvers [13] 22 [2]]. In the current work, we are particularly interested in the solution of regularized positive definite linear systems of the form
(A + Af)x = y.
which are crucial for natural gradient methods (compare to equation T). We will rely on the random- ized Nyström approximation
Anys = (A A2)* where Ω e gaxl is standard normal,
where i indicates the Mloore-Penros inverse and the rank can be chosen to control the cost and the accuracy of the approximation. This matrix provides the best positive semi-definite approximation of A whose range coincides with the range of the sketch AQ; see 25 for a thorough discussion.
3 Methods
Let us get back to the PINN example problem in Section operator Define the residuals (f); 1/VNg(Cuugrr, f(x)) with i 1...., Ng and (0))(0); = 1/ /NNoiuuu g(xg)) with j 1,..., Nen. The loss function can then be written in standard nonlinear least-squares form
2 and assume that L is a linear PDE
1 (r(0)) with r(0) L(O)= マ
(ra(0) (ran(f)
If C is nonlinear, use its linearization.

---

Algorithm SPRING for PINNs
Require: damping X, momentum μs, learning rate schedule 7/k, initial guess 0o. norm constraint C'
1: Initialize:
0 0a. 00 <-
3: for ht 1,.... K do
4: Sample new batch (#1,.... NN ....... T NNae)
5: Calculate Ta
6: Ca
D. Residual shift for SPRING; see equations Woodbury form of ENGD, equation
78
7: Φr J (JuJk X1) +c
8:
ΦE (ox + (obe-))/ p22)
Add back the shift and bias correection Sec Section 3.1 of [11
9:
0 A ox min(ns. (nt. /C/llblll)
10: end for
return Trained parameters 0
and we set N = Na + Nan and Tk = r(0%) and Ju = Jo r(0)). With this notation, it holds
G(O+)) = J Ju and V (0aa) = Jr Tk and the natural gradient scheme equation can be rewritten
0x+1 = 02 - (T Ju + AT) rk-
(4)
We again stress that, since J3 of shape (N. P) the matrix G(O)) is of shape (P,P), hence pro- hibitively large with a computational cost of ((P3) for the linear solution.
1) Leveraging Woodbury's Identity. We in the proof of the well-known Woodbury formula, to reduce the computational cost to O(N2 While these are well-known matrix identities, their relevance for natural gradient methods in the context of quantum many-body problems was only recently realized in the VMC community and proposed minimal norm stochastic reconfiguration (minSR) 4 and further refined by (36). In our notation, the push-through identity is
the push-through identity, an intermediate step
(J J& + AD Jr J (JAJH + AA) Tk-
(5)
ERPxP
ERNXN
Using the right-hand side (ENGD-W) W) instead of the left, the kernel matrix J. E RNN now lives in sample space and is thus cheap to invert. Indeed, the complexity O(N of computing the kernel matrix now dominates the computational cost, a drastic improvement over the O(Pa) cost of the original ENGD. The matrix J is the neural tangent kernel [15], and various ways of efficiently computing have been discussed in [30]. We discuss our approach the implementation details in Section 4 We stress again that using the kernel form drastically improves the scalability of ENGD, allowing network and batch sizes up to
10 and N 10 on consumer hardware while up to floating point arithmetic performing the exact scheme of equation 4 It worth noting that the O(N2P) complexity achieved by applying the Woodbury identity is still quadratic the batch size N, whereas by applying further approximations K-FAC is able to achieve linear scaling with respect to the batch size. We have not observed this quadratic scaling to be a bottleneck for realistic batch sizes.
2) Introducing Momentum. The major drawback of the kernel formulation relative to methods that work with G(Ax) directly is that it makes it difficult to aggregate curvature information over multiple training iterations, which can lead poor performance in highly stochastic settings. This problem was addressed by the introduction of the SPRING algorithm, which builds on MinSR by incorporating special type of momentum that is tailored to accelerate the convergence of natural gradient methods [1I]. To understand the SPRING algorithm, note first that the ENGD update of equation 5 is also the solution to Tikbonov-regulaizzed least-squares problem:
J (JEJH AT) - -r = arg min 1& - rxll2 + 2116112.
(6)
SPRING modifies the regularization term to incorporate the previous momentum direction Φk-1
Φv arg min 1| rx A1lø= μ
or approximations thereof, like K-FAC
5

---

Algorithm 2 GPU-Efficient Randomized Nyström Approximation
Require: A E $4 (R), target rank é, regularizer A>0 0
Ensure: (B, L) so that
1 v)).
Anys
BB",
(Anys + v A X
1: Ω randn(n, n,
Gaussian test matrix
2: Y < AQ
P eps (1)
b small shift >> embed shift in A + vl
4: Y, + Y + U
5: C k- (
6: < Y, C-1
D Cholesky solve instead of inversion b for Woodbury solve
7: R + Ba B+N
8: L + (h)
This is justified by a connection to the randomized block Kaczmarz method and yields significant acceleration relative MinSR 11 This new optimization problem has closed-form solution
Φ₂ = p J (J + (rx -JxŌx-1);
-
(8)
and the cost of evaluating this formula is only negligibly higher than for MinSR. Note that MinSR is recovered by setting μ 0. Additional insights the behavior of the iteration 8 and the function of the hyperparameters X, can be found in B 12 As a new contribution, we add bias correction of
1/ 1 - 2 to the SPRING algorithm as is customarily done in other momentum-bass optimizers
The precise algorithm is given in Algorithm !
3) Nyström on the Kernel & GPU-Efficient Implementation. We introduce randomization into ENGD and SPRING by applying randomized Nyström approximation of the N x N kernel matrix, potentially reducing the per-iteration cost even beyond the O(N2P) cost attained via Woodbury's s identity. Unfortunately, the well-known stable algorithm for applying the Nyström approximation (see 110. algorithm 2.1) requires an SVD of a N x S matrix. This cost is asymptotically negligible but nonetheless we have found that it is very significant in practice because the SVD runs extremely slowly on GPU. To ameliorate this problem, we propose GPU-efficient Nyström approximation which we summarize in Algorithm 2 Our proposal runs an order of magnitude faster GPU (see Appendix B) and enables speed-ups to be obtained even when S is only a single order of magnitude smaller than N.
The main differences relative to the standard stable algorithm are 1) we skip the QR decomposition of the test matrix Ω since random Gaussian matrices are already likely to be well-conditioned, 2) skip the SVD step and, as result, we technically return a Nyström approximation very small value of v, and 3) apply the Woodbury matrix identity (again) to reduce the cost of matrix-vector products with the inverse of the Nyström approximation.
A + v for a
We leverage this GPU-efficient randomized Nyström approximation in a sketch-and-solv frame- work [38] for example for ENGD the randomized version uses the approximation
J (J&J + Ar) = J (nys(J&Jī))
with the linear solve implemented following algorithm 2 approach can require very large sketch sizes to achieve high accuracy solutions, which has inspired the proposal of an alternative sketch-and- seaphh-annd-rrconditiooon approach which attain arbitrary accura- cics with any sketch size [10]. However, in the PINN setting the CG iterations required by such an approach involve additional differentiation passes through the operator C, which we have found nullifies any performance benefit from randomization in practice.
It is known that such a sketch-and-solve
4) Tracking the effective dimension. The accuracy of the randomized Nyström approximation depends on the effective dimension [10] of the regularized matrix A + X/. The effective dimension measures the degrees of freedom of the problem after regularization and serves as smooth quantifier of the number of eigenvalues larger than The smaller the effective dimension, the smaller the
sketch size needed to form accurate Nyström approximation and the better performance we can

---

5d Poisson (D 10065)
5d Poisson (D 10065)
10
10
I/ error
La error
10
10
10-
10-9
10"
10
10
10a
10
10
10a
10
Iteration
Time (s)
ENGD
Hessian-fres
ENGD (Woodbury)
Figure 2: 5D Poisson: Performance comparison of different optimization algorithms on a five- dimensional Poisson problem discretized with 10065 parameters, trained with tanh-activated mul- tilayer perceptron. Introducing the Woodbury matrix identity allows ENGD to take more than 30 times steps and outperform the Hessian-free optimizer by significant margin.
expect from our randomized algorithms. The effective dimension is defined as
A
dan(A) Tr
Σ
4 Experiments
To show the implications of our contributions, we present 4 experiments, each consisting of more than 1200 runs average. First, we show that using the kernel matrix improves ENGD beyond the performance of the usual baselines. Second, we show that incorporating momentum using SPRING further accelerates convergence and achieves state-of-the-art results. Moreover, SPRING removes the need for an expensive line search and performs particularly well in high dimensional settings. Third, show that our GPU-efficient Nyström approximation training for low-dimensional problems with large batch sizes; in most cases, however, it appears preferable use the exact Woodbury formula. Fourth, show that the regularized kernel matrix has an effective dimension that is only mildly smaller than the batch size N, which explains the limited benefits achieved by randomization. Importantly, our randomized algorithms still far outper- form the original ENGD method, and it is only when comparing to the new Woodbury variants that they lose their appeal.
accelerate the early phases of
Setup To demonstrate our methods, we consider a Poisson equation -Au(x) = f(x) with dif- ferent right-hand sides and boundary conditions on the unit square x E [0, 1]d for = 5, 100. For d = 5, we use as manufactured solution u(z) = D i=1 (T(,)) and right-hand side = ,²u. For = 100, we use u(x) = ||ا|| for & E a 1]100 and consequently f = - 2d. Furthermore, we consider a 4+1d Heat equation 0,u(t, x) xt,
0 for E [0, 1], [o, 1]a and bound- ary conditions u(0, x) Σi= (2n;) and u(t, = exp(-() Σi= sin(2x;), this last one with E a00 11a. . At last, we consider a 9+1d Fokker-Plande equation logarithmic space given by
d
xig(t,x) 2 Vg(t, x) x x) Aq(t,x)
((0) log(p" (0)).
where d 9,1 [o, 1] and E R.. practice, x E [-5,5] The solution g given by q log(p") with p" (1, ~ N(0, exp(-t) + (1-exp( (oo(-t))2 2)) For further results see Appendix A
Implementation We use the same sequential architecture consisting of four hidden layers and single output for all problems, varying the width of the layers as function of d. Moreover, we imple- ment the model using Taylor-mode forward differentiation [2), and use Jacobian vector products to optimize kernel creation and kernel vector products [15). To implement matrix-free alternative to ENGD, we adopt the Hessian-fres free optimization method 23 which applies truncated conjugate gradient iterations and computes exact Gramian- vector products to improve gradient conditioning.
7

---

5d Poisson (D 10065)
100d Poisson (D 1325057)
10a
100
10-1
La errr
10-1
10-3
10
10
100
10
10a
10a
100
10
10a
10a
Time [s]
Time [s]
4d Heat (D 116865)
9+1d log Fokker Planck (D 118145)
10"
10a
L error
10
10
10
10-2
10-3
10P
10a
10a
10a
10
10a
10a
10a
Time [s]
Time [s]
ENGD (Woodhury)
SPRING
ENGD-W W (Iine Search)
KFAC
Figure 3: Performance of ENGD and SPRING on four problems: 5d and 100d Poisson, 4+1d Heat, and 9+1d log-Fokker-Plancks SPRING achieves 1 errors not previously seen in the high dimensional settings, the 100d Poisson and 9+1d gg--Fokker-Plande problems. In the latter, the error an order of magnitude lower that the previous state-of-the-an KFAC.
We evaluate two other baselines: SGD with tuned momentum and learning rate, and Adam with tuned learning rate. Hyperparameter tuning is performed using Weights & Biases 47]. as detailed in Appendix A Our tuning strategy involves an initial exploratory stage with 50 trials over broad parameter ranges, followed by refinement stage with another trials narrowed regions in- terest. Performance is evaluated based on the lowest Le-error in validation set with known PDE solution. All experiments are run on a uniform hardware setup, an RTX 6000 GPU cluster (24 GiB memory) using double precision, and each optimizer is given an equal compute time budget on the same fixed PINN task. The complete ranges of byperparameters, the selected configurations and the dynamics of training in terms of the iteration count are available in Appendix A
1) Introducing Woodbury into natural gradient methods for PINNs. We first reproduce the 5-dimensional Poisson problem of Dangel, Müller, and Zeinhofer (6. with batch size of 3500. We compare the baselines against ENGD-W (see Figure D. and verify that the exact ENGD already attains top-tier solution accuracy. Incorporating the Woodbury identity then increases the iteration speed by more than a factor 30 and surpasses all other approaches, including the original ENGD and Hessian-free optimization, in both efficiency and accuracy.
Woodbury is the right way to implement second-order methods such as ENGD.
2) Incorporating momentum further accelerate ENGD. We next incorporate momentum using the SPRING algorithm. Figure 3 shows that momentum increases the final accuracy in all cases, with a large benefit in the high-dmensional setting, i.e. the 100D Poisson problem and the 9+1d Nog-Fokker-Planct problem. It is notable that the SPRING algorithm does not use line search and is able to outperform ENGD both with and without its line search. The final accuracy attained by both ENGD with Woodbury and SPRING surpasses even that attained by (6] using the more complicated KFAC algorithm. This is accentuated with SPRING attaining an T.--error an order of magnitude lower than KFAC in the 9+1d log- Soee--ooke-eeee Planck problem.

---

N 1000
N 5000
N 10000
L2 emor
10-1
10-3
10-5
10
10
10
10
10
10
10
10
10
10
10'
10
10
10
10
Time [8]
Time [s]
Time [s]
ENGD (Woodhury)
ENGD (Nystrom)
Figure 4: Effect of Nyström ENGD for the 5D Poisson equation.
SPRING yields state-of-the-art performance without the complexity of KFAC.
3) Randomization for large batch sizes. We have shown that adding momentum to ENGD-W with tuned learning rate is the best strategy to optimize PINNs. When the batch size is very large, even the computation of the kernel matrix can be quite expensive, motivating the use of randomized algorithms reduce the cost. We explore the of randomized Nyström approximation celerate the computation and inversion of the regularized kernel matrix J + XZ. In the first set of experiments, shown in Figure 4 we compare ENGD-W with two randomized variants, one us- ing the standard stable Nyström approximation and the other using our new GPU-efficient variant, Algorithm 2 We test batch sizes of 1000, 10 000, and 50000, all under our standard line-search procedure and all with sketch size of 10% of the batch size. The 10% batch size is chosen because larger batch sizes, even the GPU-efficient algorithm cannot provide substantial cost reductions. We find that randomization can accelerate the early phase of the training, especially when the batch size large, but ultimately exact computations are needed to reach high accuracy results.
In the second set of experiments, shown in Figure 5 we compare SPRING against its two corre- sponding randomized variants randomized algorithms is approximately equal to or worse than the performance of SPRING with exact computations. The main reason that we do not observe any speed-ups from randomization here is that for nigh-dimensional problems, the cost of the training is increasingly dominated by the cost of differentiating through the operator L. which scales linearly with the problem dimension. Thus, the acceleration of the kernel computation and inversion becomes less relevant. For the 5d problem, we find a speedup 2 for a sketch size of 10% of N, and no speedup for values above 25% of N. Given the cost of the matvecs and the Laplacian operator, this is to be expected.
100D Poisson problem. In this case, the performance of the
Randomization accelerates the early phases of training for low-dimensional problems with large batches. It remains an open problem to extend this acceleration broader
scenarios.
N = 100
N 500
N 1000
In error
10
10
10
10a
10
10
10a
10a
10
101
10a
10a
10a
10
10a
10a
10a
Time [s]
Time
Time
[s]
SPRING
SPRING (Nystrom)
Figure 5: Effect of Nyström on SPRING for the 100D Poisson equation.

---

ENGD (Wondhury)
SPRING
0.8
der/N /
dat /N
0.6
0.8
0.4
0.7
0.2
10
10a
102
10a
10
10a
10a
Steps
Steps
(a) Effective dimension of ENGD's regularized ker- nel matrix in the Poisson experiment relative to the batch size. N 3500
(b) Effective dimensions of SPRING's regularized kernel matrix the 100D Poisson experiment rela- tive to the batch size. N 150
4) Effective dimension of the regularized kernel matrix. To understand the limited benefits of randomization, we track the effective dimension of the kernel matrix for both ENGD-W in the 5D problem and SPRING the 100D problem, with results Figure 6a and Figure [bb respectively. The effective dimension plateaus at more than 50% of the batch size, which explains why random- ization with sketch size of 10% results in a loss of accuracy. Unfortunately, the gap between the effective dimension and the batch size N is not large enough to squeeze out significant performance gains without losing accuracy.
Randomization suffers because the sketch sizes needed to attain significant cost savings are too small accurately approximate the kernel matrix.
Our findings contrast those of McKay et al. 226 which reported benefits from using randomization in original version of ENGD. We compare our randomized algorithms to the much stronger baseline of ENGD-W, and it is exactly the introduction of the Woodbury identity that nullifies the benefits of randomization by shifting the problem from the high dimensional parameter space to the much lower dimensional sample space, leaving little room for further dimensionality reduction.
Conclusion
We have significantly improved energy natural gradient descent for PINNs by: 1) Woodbury's identity accelerate the inversion of the kernel matrix, thus slashing per-iteration complexity, 2) SPRING, momentum scheme that accelerates the convergence of natural gradient methods, and 3) GPU-efficient Nyström sketch-and-solve approach approximate the kernel matrix for further cost reductions in some settings. We demonstrate our methods across four settings: 5d and 100d Poisson equations, 4+ld Heat equation and 1d Fokker-Plaarl equation with batch sizes up to 10 000. In the 5d Poisson case, our methods achieve the same sub-10 Z as ENGD up to 75x faster, while in the 100d case. SPRING outperforms all previously seen optimizers. In the 4+ Id Heat case our methods perform competitively with previous state of the art methods like KFAC, and the 9+1d log- Pog-Fokker-Plann Plankk problem, SPRING achieves 1/2eerrrrs an order magnitude lower than previous methods. Regarding randomization, we find promising results for low-dimensional problems with large batch sizes, and we identify barriers to achieving acceleration more generally.
Limitations & future directions Despite promising results in low-dimensional. large batch set- tings, our randomized approaches exhibit diminishing returns in the late optimization stages and fail to scale to high-dimensional problems. Future works can explore alternative approaches to randomization, but will need to contend with the peculiarities of the PINN setting in which matrix- vector products with the kernel are very expensive. Moreover, our current analysis focuses on fixed Nyström rank, leaving open questions about how sketch dimension and adaptive rank selection af- fect performance across diverse tasks. Future work should also explore how to adaptively set the hyperparameters of the ENGD-W and SPRING algorithms to yield fast, black-box optimizers.
10

---

Acknowledgments and Disclosure of Funding
M.Z. acknowledges support from ETH Postdoctoral Fellowship for the project "Reliable, Effi- cient, and Scalable Methods for Scientific Machine Learning". G.G. acknowledges support from the U.S. Department Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE- SC00231 12. A.G-C thanks Jeffrey Ren for helpful observations on the optimizer's setting used in plot generation. Resources used in preparing this research were provided, in part, by Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
com-ccontrraint
This report was prepared as an account of work sponsored by an agency of the United States Government. Neither the United States Government nor any agency thereof, nor any of their employees, makes any warranty, express or implied, assumes any legal liability or responsibil- ity for the accuracy, completeness, or usefulness of any information, apparatus, product, process disclosed, or represents that its use would not infringe privately owned rights. Reference herein to any specific commercial product, process, or service by trade name, trademark, manufacturer, or otherwise does not necessarily constitute or imply its endorsement, recommendation, or favor- ing by the United States Government or any agency thereof. The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Government or any agency thereof.
Disclaimer:
References
[1] Becca, F. and Sorella, S. Quantum Monte Carlo approaches for correlated systems. Cambridge University Press. 2017.
[2] Bettencourt, Johnson, M. J., and Duvenaud, D. Taylor-mode automatic differentiation for higher-order derivatives in JAX. In Program Transformations for ML Workshop at NeurlPS 2019, 2019. URL https::s::/ //pperrvveeeeeeee nett/forurFiddskxE333FPP
[3] Carleo, G. and Troyer, M. Solving the quantum many-body problem with artificial neural networks. Science, 35566325),602-666, 2017.
[4] Chen, A. and Heyl, M. Efficient optimization of deep neural quantum states toward machine precision, arXiv preprint urXiv:2302.019411 2023.
[5] Dangel, F., Tatzel, L., and Hennig, P. ViVIT: Curvature access through the generalized gauss- newton's 's low-rank structure. Transactions on Machine Learning Research (TMLR), 2022.
[6] Dangel, F., Müller, J., and Zeinhofer, M. Kronecker-factored approximate curvature for inysics-informed neural networks. arXiv preprint arXiv:2405. 15603, 2024.
[7] Dissanayake, M. W. M. G. and Phan-Thien, N. Neural-network--usedd approximations for solving partial differential equations. Communications in Numerical Methods in Engineering, 10(3):195 201, March 1994. ISSN 1099-0887. doi: 1..1002/cm..66400030000 1640100303. URL http:
//dx. didx.doi.og org. 10. 1002/ cnm. 1000/c0m.1640000330000000
[8] Epperly, E. N., Meier, M., and Nakatsukasa, Y. Fast randomized least-squares solvers can be just as accurate and stable as classical direct solvers. arXiv preprint arXiv:2406. 03468, 2024.
[9] Foulkes, W. M., Mitas, L., Needs, R., and Rajagopal, Quantum Monte Carlo simulations of solids. Reviews of Modern Physics, 73(1):33, 2001.
[10] Frangella, Z., Tropp, J. A., and Udell, M. Randomized Nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2):718-752, 2023.
[11] Goldshlager, G., Abrahamsen, N., and Lin, L. Kaezmarz-inspired approach to accelerate the optimization of neural network wavefunctions. Jourmal of Computational Physics, 516: 113351, 2024.
[12] Goldshlager, G., Hu, J., and Lin, I.. Worth their weight: Randomized and regularized block Kaczmarz algorithms without preprocessing. arXiv preprint 22XXx:2502200882 2025.
11

---

[13] Halko, N., Martinsson, P.-.,, and Tropp. J. A. Finding structure with randomness: Proba- bilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53(2): 217-288, 2011.
[14] Hermann, J., Spencer, J., Choo, K., Mezzacapo, A., Foulkes, W. M. C., Pfau, D., Carleo, G., and Noé, F. Ab initio quantum chemistry with neural-network wavefunctions. Nature Reviews Chemistry, 7(10);692-709 2023.
[15] Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In Advances Neural Information Processing Systems, 2018
[16] Jnini, A., Vella, F., and Zeinhofer, M. Causs-Newton natural gradient descent for physics- informed computational fluid dynamics. arXiv preprint arXiv:2402.10680 2024.
[17] Kingma, D. and Ba, J. Adam: A method for stochastic optimization. International Confer- ence on Learning Representations, 2015.
[18] Kiyani, E., Shukla, K., Urbán, F., Darbon, J., and Karniadakis, G. E. Which optimizer works best for thysics-informed neural networks and Kolmogorov-Amold networks?, 2025. URL
wttp:xii..rgg/bb///5//00 16377
[19] 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 2--26666-66600, 2021.
[20] Lagaris, I., Likas, A., and Fotiadis, D. Artificial neural networks for solving ordinary and par- tial differential equations. IEEE Transactions on Neural Networks, 95):9987-100, 1998. ISSN 1045-9227. doi: 10.1 109/72.71177 URL http:ttp//dae doi org/ 1109/72. 11097/22..10//
and Chen, J. Ab initio calculation of real solids via neural network ansatz. Nature [21] Communications, 13(1):7895, 2022.
[22] Lin, J., Goldshlager, G., and Lin, L. Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation. Journal of Computational Physics, 474:111765,, 2023.
[23] Martens, J. Deep learning via Hessian-free optimization. In International Conference on Machine Learning (ICML), 2010.
[24] Martens, and Grosse, R. Optimizing neural networks with Kronecker-factored approximate curvature. In International Conference on Machine Learning (ICML), 2015.
[25] Martinsson, P.-G. and Tropp. J. A. Randomized numerical linear algebra: Foundations and algorithms. Acta Numerica, 29:403-572, 2020.
[26] Mckay, M. B., Kaur, A., Greif, C., and Wetton, B. Near-optimal sketchy natural gradients for physics-informet neural networks. Forty-second International Conference on Machine Learning, 2025.
[27] Müller, and Zeinhofer, M. Achieving high accuracy with PINNs via energy natural gradient descent. In International Conference Machine Learning, pp. 25471-25485 PMLR, 2023.
[28] Müller, J. and Zeinhofer, M. Position: Optimization in SciML should employ the function space geometry. In Forty-first International Conference on Machine Learning, 2024.
[29] Murray, R., Demmel, J., Mahoney, M. W., Erichson, N. B., Melnichenko, M., Malik, O. A., Grigori, L., Luszczek, P., Dereziáski, M., Lopes, M. E., et al. Randomized numerical linear algebra: A perspective on the field with an eye software. arXiv preprint arX1v:2302.11474, 2023.
[30] Novak, R., Sohl-Dickstein, J., and Schoenholz, S. S. Fast finite width neural tangent kerne.. In International Conference on Machine Learning, pp. 77018-17044 PMLR, 2022.
[31] Pearlmutter, B. A. Fast exact multiplication by the Hessian. Neural Computation, 1994.
12

---

[32] Pfau, D., Spencer, J. S., Matthews, A. G., and Foulkes, W. M. C. Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical review research, (3):033429, 2020.
[33] Raissi, M. Perdikaris, P., and Kariiadaiis, G. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial dif- ferential equations. Journal Computational Physics, 378:686-707, February 2019. ISSN 0021-9991. doi: jep..16/.. 10.044. URL http::/dae doi org/ 1016///...../j...ppp...p..
2018.10.04 14.
[34] Rathore, P., Lei, W., Frangella, Z., Lu, L., and Udell, M. Challenges in training PINNs: loss landscape perspective. arXiv preprini M.Xiv:2402.01666 2024.
[35] Ren, Y. and Goldfarb, D. Efficient subsampled Gauss-Newton and natural gradient methods for training neural networks. arXiv, 2019.
[36] Rende, R., Viteritti, L. L.,, Bardone, L.,, Becca, F., and Goldt, S. A simple linear algebra identity optimize large-scale neural network quantum states. Communications Physics, (1):260, 2024.
[37] Rokhlin, V. and Tygent M. A fast randomized algorithm for overdetermined linear least- squares regression. Proceedings of the National Academy of Sciences, 135136)::3212-1321,,, 2008.
[38] Sarlos, T. Improved approximation algorithms for large matrices via random projections. In 2006 47th annual IEEE symposium on foundations of computer science (FOCS'06), pp. 143- 152. IEEE, 2006.
[39] Schätzle, Z., Szabó, P. B., Mezera, M., Hermann, J., and Noé, F. DeepQMC: An open-source software suite for variational optimization deep-learning molecular wave functions. The Journal of Chemical Physics, 159(9), 2023.
[40] Schraudolph. N. N. Fast curvature matrix- vector products for second-order gradient descent.
Neural Computation, 2002.
[41] Singh, S. P. and Alistarh, D. Woodfisher: Efficient second-order approximation for neural net- work compression. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[42] Sorella, S. Generalized Lanczos algorithm for variational quantum Monte Carlo. Physical
Review B, 64(2):024512, 2001.
[43] Tatzel, L., Hennig, P., and Schneider, F. Late-phase second-order training. In Has it Trained Yer? NeurlPS 2022 Workshop, 2022. URI JRRl:/tpeee//ppnnreeeeeeeeee/eee wtwp:/wwwwwiiieteennneee.aa
C3hL1sb2511
[44] Wang, S., Teng, Y., and Perdikaris, P. Understanding and mitigating gradient flow pathologies in physicc-iiniiiiee informed neural networks. SIAM Journal on Scientific Computing, 43(5):A3355- A3081, 2021.
[45] Wang, S., Yu, X., and Perdikaris, P. When and why PINNs fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449:1 10768, 2022.
[46] Wang, S., Bhartari, A. K., Li, B., and Perdikaris, P. Gradient alignment in physics-informed neural networks: second-order optimization perspective, 2025. URL ltttss:/arxiv. 00604
[47] Weights and Biases.
Experiment tracking with weights and biases. htttp::/www.wanddee ai/ 2020. Software available from wandb.ai.
[48] Williams, C. and Seeger, M. Using the Nyström method to speed up kernel machines. Advances in neural information processing systems, 13, 2000.
[49] Yang, M., Xu, D., Wen, Z., Chen, M., and Xu, P. Sketch-based empirical natural gradient methods for deep learning. Journal of Scientific Computing, 92(3):94, 2022.
13

---

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 contributions bullet points in Section 1 and provide references to the parts the paper that present them.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
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 the work performed by the authors?
Answer: [Yes]
Justification:
We provide paragraph limitations in Section 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 discussed in the paper.
The authors are encouraged to create 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-specificatton, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because 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 ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed 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 complete (and correct) proof?
Answer: [NA]
14

---

Justification:
We provide derivations for the proposed methods, but our work does not present any rigorous mathematical statements.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
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 provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the of the paper should be comple- mented by formal proof' 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 conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide full description of the experimental protocol in Appendix [A and details for each individual experiment in the same appendix, including all syper-parameter search spaces and hyper-parameters of the best runs.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can 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., g., in the case of a large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which thors 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 have some path reproducing or verifying the results.
15

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We will open-source implementations, well the code fully repro- duce all experiments and the original data presented in this work.
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/
Sablit//uuides/Cddeetbmii wwwpaaumissssooooooooooo cessinnooicyyl for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.coo httppp/cccppullcguuueeeeeeeeee Sudeeeummm ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how 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 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 possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Ys]]
Justification:
We list all the details in Appendix 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 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 appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We do not show any mean and standard deviations over different model ini- tialization for training curves, however, we use a thorough tuning protocol avoid artifacts from insufficient hyperparameter tuning and conduct experiments with random searches to ensure the consistency of our results.
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, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

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 the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably 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 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 the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We describe the hardware in Section and Appendix A tion time can be inferred from the description of the tuning protocol in combination with the assigned time budget per run.
The total computa-
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well 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 t make it into the paper).
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:
Answer:
[Yes]
Justification:
We have read the Code of Ethics and believe our submission aligns with
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 special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [No]
Justification: We provide collection of improvements for the accurate training of physics- informed neural networks. This classical numerical problem in applied mathematics. and we do not believe it has any immediate negative societal impact.
17

---

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 spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does release any model data, thus not applicable. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: Our work 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 used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
18

---

For scraped data from a particular source (e.g., website), the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, patasets.paperswittte paperswithcode www.comm/aaa.oc condattseet has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: Our work 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/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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
For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Question:
Answer: [NA]
Justification: Our work 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 contri- bution 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, cura- tion, or other labor should paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed 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?
Question:
Answer: [NA]
Justification: Our work does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equiva- lent) 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 include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Our work does not make use LLMs as a component any way.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy http::///meerrppssccc///oeeerrrs for what should or should not be described.
20

---

A Experimental Details and Additional Results
A.1 Hyper-Parameter Tuning Protocol
We tune the following optimizer hyper-parameters and otherwise use the PyTorrc 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 no), maximum number of CG iterations
ENGD: damping, factor of the exponential moving average applied to the Gramian, initial-
ization of the Gramian (zero or identity matrix)
ENGD (Woodbury): damping, learning rate (when fixed)
SPRING: damping, momentum, learning rate (when fixed)
Randomized: damping, learning rate (when fixed), sketch size
We use random search from Weights & Biases to determine the saper-parameter, where run is executed in double precision and allowed to run for given time budget, and we rank runs by the final 7.2 error on fixed evaluation data set. To compare, the hardware used was RTX 6000 GPUs with 24 GiB RAM. The first state consists of relatively wide search space and limit to 50 runs. The second stage is narrowed down to smaller space based on the first stage, then re-run for another 50 runs. We will release the details all hyper-parameter search spaces, as well as hyper-parameters for the best runs in our implementation.
A.2 5d Poisson Equation
Setup We consider five-dimensional Poisson equation -A)(z) > cos(TT;;) the five-dimensional unit square a E [0, 1]s with cosine sum right-hand side and boundary conditions u(x) Σ (T2;)) for alo, 115. We sample training batches of size N₂ 3000, 500 and evaluate the 2 eтror (x) Σ 1 cos(TT;; All optimizers sample
separate set of 30 000 data points using the known solution
training batch each iteration, and each run is limited to 7000s. We use an MLP five-layer architecture whose linear layers are Tanh-activated except for the final one: 5 - 64 + 64 48 48 + 1 MLP with D 10065 trainable parameters. ENGD-W and SPRING here make use of the inherited ENGD line search. Figure visualizes the results.
Best run details
The runs shown in Figure correspond to the following hyper-parameters.
SGD: learning rate: 2.895 360 x 10-3, momentum: 3 x 10
Adam: learning rate: 2.808 451 10-4
Hessian-free: curvature matrix: GGN, initial damping: 1 x 10-1, constant damping: no, maximum CG iterations: 350
ENGD: damping: 1 10- exponential moving average: 0, initialize Gramian to identity: yes
ENGD-W: damping: 3.173 10-12
SPRING: damping: 2.086 287 10,10; momentum: 3.115 10-1
Search space details The runs shown in Figure were determined to be the best via search with
approximately 50 runs on the following search spaces which were obtained by refining an initially wider search (l denotes a uniform, and LU log-uniform distribution):
SGD: learning rate: Lu((1 x 10 1x 10 2I), momentum: U([0,3 10 -..0-.6.. -,, 10 10 10-1})
AAdm: learning rate: Lu([1 x 10-1;5 10-11)
21

---

5d Poisson (D) 10065)
5d Poisson (D) 10065)
10a
10
L2 etror
Loss
10
10-7
10
10
10a
10
10a
10
10
10
10
10
10
10
10a
Iteration
Iteration
5d Poisson (D = 10065)
5d Poisson (D = 10065)
10a
10
La error
Loss
10
10-7
10
10
10
10a
10a
10
10a
10
10a
10a
10'
Time (8)
Time (s)
SGD
Adam
ENGD
Hessian-ffee
ENGD (Woodbury)
SPRING
Figure 7: Training loss and evaluation 2 error for learning the solution 5d Poisson equation over time and steps.
Hessian-free: curvature matrix: 4((GG))), initial damping: 4{{100, 50, 10, 1,5 x 10-1.1 x 10-1,5 5 x 10-2)), constant damping: U((no)), maximum CG iterations: U({100, 150, 200, 250, 300, 350))
ENGD: damping: U({1x 10 111100 1110011110 10 1x100111 1x100-11 0}), exponential
moving average: u(00,3 3 10-1,6 10-1,9 10-1)), initialize Gramian to identity: u([no. yes})
ENGD-W: damping: [U((1 10-7;1)
SPRING: damping: [U[(] 10-10; 10= momentum: U((0.0; 0.999))
A.2.1 Fixed learning rate
We repeat the previous experiments, now adding a search space of [[(( x 10-1; x 10-41) for the learning rate, see Figure We note that this change meant adjusting the search space for SPRING's momentum to (0.8 0.999]). We find that the best parameters are:
ENGD-W: damping: 6.804 474 10
learning rate: 5.2289 10-2
SPRING: damping: 6.811 585 10-10, momentum: 8.269 66 10-1, learning rate:
6.3502 x 10-2
A.2.2
Large batches
We now repeat the experiment using batch sizes of 1000, 5000, and 10 000 with the line search, to test the randomized approach, setting the sketch size to 10% of N. We can visualize the results in Figure
A.3
10d Poisson Equation with line search
Setup We consider 10-dimensional Poisson equation -u((z) 0 on the 10-dimensional unit square & [0, 1 with zero right-hand side and harmonic mixed second order polynomial boundary
conditions u(x)) = Ex; 221-1 T2z for æ 1001 We sample training batches of size No =
22

---

5d Poisson (D 10065)
5d Poisson (D 10065)
107
10a
L error
Loss
10
10"
10
10
10a
10a
10a
10
10a
10a
Iteration
Iteration
5d Poisson (D 10065)
5d Poisson (D 10065)
10
10a
L error
Loss
10
10
10
10a
10a
10
10'
10a
10a
10"
Time
Time
ENGD (Woodbury)
SPRING
Training loss and evaluation L error for learning the solution 5d Poisson equation over time and steps with fixed learning rate.
Figure 8:
3000, Nan 1000 and evaluate L error on a separate set 30 000 data points using the known solution (x) = Σ4/ I2i- Both optimizers sample a new training batch each iteration, and each run limited to 7000 We use a 10 - 256 - 256 128 + 128 1 MLP with D = 118 145 MLP whose linear layers are Tanh-activated except for the final one. Given the poor performance of SGD, Adam, and the Hessian-free, we no longer run them in this more complicated problem. Furthermore, the traditional ENGD runs out of memory for networks of this size. Figure visualizes the results.
Best run details
The runs shown in Figure correspond to the following hyper- hyper- parameters:
-ENGD-: damping: 3.9 10-7
SPRING: damping: 1.7 10-7 momentum: 9.053 28 10
Search space details The runs shown in Figure 11 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 (l denotes uniform, and LU log-uniform distribution):
ENGD-W: damping: LU((1 10)
SPRING: damping: CU((1 10-10; x 10-3)); momentum: (2((606009999)
A.3.1 Fixed learning rate
We repeat the previous experiments, now adding a search space of LU( x 10-1; x 10 10-4)) for the learning rate, see Figure [2] We find that the best parameters are:
ENGD-W: damping: 1.579 10-5; learning rate: 8.7024 x 10
SPRING: damping: 4.98 x 10-5; momentum: 9.6764 x 10-1; learning rate: 6.034 67 x
10-2
23

---

N 1000
N 5000
N 10000
Loss
10
10a
10
10
100
10
10
10
Iteration
Iteration
Iteration
1000
5000
10000
La error
10-1
1071
10
10a
10a
10a
10
10a
10a
10*
10a
10
10a
10
Iteration
Iteration
Iteration
N = 1000
N 5000
10000
Loss
10-2
10" 7
100
10
10a
10
10
10
10a
10a
10a
10a
10a
Time [s]
Time [s]
Time [s]
N 1000
N 5000
N 10000
La error
10
10-1
10
10
10
10a
10
10
10a
10a
10a
10
10a
10
10a
10a
10
Time
Time
[s]
Time [s]
ENGD (Woalbury)
ENGD (Nyurom)
SPRING
SPRING (Nystrom)
Figure 9: Training loss and evaluation 2 error for learning the solution ENGD-W a 5d Poisson equation over time and steps with large batch sizes and randomization.
A.4 100-d Poisson Equation with line search
Setup Here, we consider a 100d Poisson equations -u((x) f(x) with zero right-hand side f(x) = 0, harmonic mixed second order polynomial boundary conditions u(x) = Σ T2i-1
for æ E 0[o, 1]d, and known solution (x)) Σ -2:2 142;, We assign each run a budget of 10 s. We tune the oppimizer-hyperpaaaneee
described in Appendix A. using random search. We use 100 + 768 + 768 - 512 > 512 + MLP with layers are Tanb-activated except for the final one. This architecture is again too large for ENGD to optimize. Figure 13 visualizes the results.
325 057 MLP whose linear
Best run details
The runs shown in Figure correspond to the following hyper- parameters:
LNNGDW: damping: 4.7772 10-3
SPRING: damping: 3.0106 x 10-2 momentum: 6.763 35 x 10-1
Search space details The runs shown in Figure 13 were determined to be the best via random search on the following search spaces which each optimizer given approximately the same total computational time (ll denotes uniform, and Lll log-uniform distribution):
ENGD-W: damping: LU((1 x 10-10; 10-1])
SPRING: damping: LU((1 10-10; 10 1 10 10-31); momentum: [0.6, (Z.(69909999))

---

N 1000
N 5000
N 10000
107
Loss
10
10-9
10
10a
10a
10a
10
10a
10
Iteration
Iteration
Iteration
1000
5000
10000
10
La error
10
10
10a
10a
10a
10
10a
10a
10
10a
10
10a
10a
Iteration
Iteration
Iteration
N 1000
N 5000
N 10000
107
Loss
10a
10-3
100
10
10a
10
10a
10a
10a
10
10a
10a
10
Time [s]
Time [s]
Time [s]
N 1000
N 5000
N 10000
10
La error
10" 8
10
10
10a
10a
10
10
10a
10a
10a
10
109
10
10a
10a
10
Time
Time [s]
Time
[s]
ENGD (Woalbury)
ENGD (Nyurom)
SPRING
SPRING (Nystrom)
Training loss and evaluation I error for learning the solution to 5d Poisson equation over time and steps with large batch sized and fixed learning rate.
Figure 10:
A.4.1 Fixed learning rate
We repeat the previous experiments, now adding search space of cu(( 10-1; 10-41) for the learning rate, see Figure 12 We find that the best parameters are:
ENGD-W: damping: 6.233 10-7; learning rate: 9.118 10
SPRING: damping: 3.0116 x 10-2; momentum: 9.8386 10-1; leeannnng rate: 9.2362 x
10-2
A.4.2 Large batches
We now repeat the experiment using batch sizes of 1000, 5000. and 10 000 with the line search, in order to test the randomized approach, setting the sketch size to 10% and fixed the learning rate with the previously introduced search space. We visualize the results in Figure 9
A.S 4+1d Heat equation
Setup We consider a 4+1- --dimemsional heat equation O.u(t, x) (rt,u(t,o the four-dimensional unit square and unit time interval, v, E [0, 1j4 [0. The equation has
with x 1 on
spatial boundary conditions u(t, x) exp(-t) Ei= sin(2x;) fort, x E [0, 1] x 0 1]4 throughout
time, and initial value conditions u(0, æ)) Ei= sin(2r;) for [0, 1Ja We sample training
25

---

10d Poisson (D 118145)
10d Poisson (7) 118145)
10a
10
La error
Loss
10-4
10
10
10"
10
10a
10"
10
10a
Iteration
Iteration
10d Poisson (D = 118145)
10d Poisson (D = 118145)
1000
10-2
L error
Loss
10-1
10
10a
10a
10a
10a
Time (8)
Time
ENGD (Woodbury)
SPRING
Figure 11: Training loss and evaluation L error for learning the solution 10d Poisson equation over time and steps.
10d Poisson (D 1181455
10d Poisson (D 1181455
10
10
L/ error
Loss
10
10
10a
10
10
10a
10
10
10a
10a
10a
10a
Iteration
Iteration
10d Poisson (D 118145)
10d Poisson (D 1181455
10
10
error
Loss
10
10
10a
10a
10a
10
Time (s)
Time (s)
ENGD (Woodbary)
SPRING
Figure 12:
Training loss and evaluation L error for learning the solution 10d Poisson equation over time and steps with fixed learning rate.
batches of size Ng = 3000, Nao = 500 (Nan/2 points for the initial value and spatial boundary conditions each) and evaluate the 2 error on separate set of 30000 data points using the known solution u. (t, exp(-t) Ei= sin(2xx;) We sample new training batch each iteration. Each
26

---

100d Poisson (D) 1325057)
100d Poisson (D 1325057)
10'
La error
Loss
10-1
10
10
10
10a
10a
10a
10
10a
10a
Iteration
Iteration
100d Poisson (D = 1325057)
100d Poisson (D = 1322057)
10
La error
Loss
10-1
10-3
10
10a
10a
10'
10'
10
10a
10'
Time (8)
Time
ENGD (Woodbury)
SPRING
Figure 13: Training loss and evaluation L error for learning the solution high-dimensional Pois- son equations over time and steps using random search.
100d Poisson (D = 1325057)
100d Poisson (D 1325057)
10a
error
Loss
10-2
La
10a
10
10a
10a
10a
10a
10
10a
10a
10a
Iteration
Iteration
100d Poisson (D 1325057)
100d Poisson (D 1325057)
10
La error
Loss
10-2
10
100
10a
10a
10
10*
10
10a
10a
10a
10
Time (8)
Time (s)
ENGD (Woodbary)
SPRING
Figure 14: Training loss and evaluation 2 error for learning the solution a 100d Poisson equation over time and steps with fixed learning rate.
run is limited 3000 We fayer MLP architecture whose linear layers Tanh-activated
except for the final one: 5 + 256 + 256 + 128 + 128 + with = 116864 trainable weights.
Best run details
The runs shown in Figure 16 correspond to the following
27

---

N = 100
N 500
N 1000
10a
Loss
10-3 -
10
10*
10
10
10a
10
10
Iteration
Iteration
Iteration
100
500
La error
10"
10-1
10-2
10a
10
10a
10
10
10a
10
10a
10a
10a
10
10a
10a
Iteration
Iteration
Iteration
= 100
N = 500
1000
10"
Loss
10-3 -3
100
10
10
10a
10
10a
10a
10
10a
10
Time [s]
Time [s]
Time [s]
100
N 500
N 1000
Lo error
10
10
10
10
10
10a
10a
10
10a
10a
10a
10a
10
10
10
10a
10a
Time [s]
Time [s]
Time
SPRING
SPRING (Nystrom)
Figure 15: Training loss and evaluation I,2 error for learning the solution to a 100d Poisson equation over time and steps with large batch sized and fixed learning rate.
ENGD-W: damping: 1.933 629 x 10-7
SPRING: damping:
1.081 x 10-7, momentum: 8.500 992 10
Search space details The runs shown in Figure 16 were determined to be the best via random search on the following search spaces which each optimizer given approximately the same total computational time (Ul denotes uniform, and LU log-uniform distribution):
ENGD-W: damping: LU((1 10-10; 10-'l)
SPRING: damping: LU((1 10-19; 10-3)); momentum: [2((006;00900)
A.5.1 Fixed learning rate
We repeat the previous experiments, adding search space of CU((1 10 10-"I) for the learning rate, see Figure 17 We find that the best parameters are:
ENGD-W: damping: 1.139 970 x 10-7, learning rate: 9.939 225 x 10-2
SPRING: damping: 2.081 775 x 10-7 momentum: 9.078 456 x 10-1, learning rate: 8.663 887 x 10-2
28

---

4d Heat (D 116865)
4d Heat (D 116865)
10
10
La error
Loss
10-5
10-9
10
10
10a
10"
10
10a
Iteration
Iteration
4d Heat (D 116865)
4d Heat (D = 116865)
10a
Loss
10
La ervor
10
10
10'
10
10a
10
10a
10a
Time (8)
Time
ENGD (Woodbury)
SPRING
Figure 16: Training loss and evaluation L error for learning the solution the 4+1dd Heat equations over time and steps using random search.
5d Heat (D 117121)
Heat (D 1171211
10a
10
error
Loss
10"4
10
10-9
10"
10
10a
10a
10
10a
Iteration
Iteration
5d Heal (D 117121)
Sd Heal (D 117121)
10a
10
Le error
Loss
10-1
10
10
10a
10a
10a
10a
10a
10a
Time (s)
Time (8)
ENGD (Woodbery)
SPRING
Figure 17: Training loss and evaluation L error for learning the solution to a 4 + 1d Heat equation over time and steps with fixed learning rate.
29

---

A.6 9+1d Fokker-Planck equation in logarithmic space
Setup For a given drift μ [0, x R - R and diffusion coefficient [0, 1] + Raxd the
Fokker-Plane Planck equation with initial probability density po is given by
1 dip (V. "Trroo T v2p) 0, p(0) pu.
which is posed [0, 1] R. Note that t() is probaility density Rd for all [0, 1]. We
transform the above equation into logarithmic space via = logp. Then, g solves
1 22100 T Vall2 Tr(oo V2g) =0 ((0) = log po: 0.g + (V, + (Vg.u) -
For our experiment, we set (t, = 2 and /2T E Rdxd and replace the unbounded domain by [0, 1] |-5, 5J Finally, our solution
log p. where p. (O exp(-()f + (1 - exp( (t)) Our loss includes the PDE residual and its initial conditions. We use a tanh- activated five-layer MLP: 10 - 256 + 256 128 128 + and use batch sizes of Nę 3000
and Nao = 1000. Each run has an allocation time budget of 6000 s.
9d Log-sfokkerplaanck--orrropp (D 118145)
9d Logg-fokkerrpllnc--ioorrppi (D 118145)
10a
10a
Loss
10
La
10-3
10a
10
10a
10a
10
10a
Iteration
Iteration
9d Log- horg-fookerrpppp plancctoo (D 118145)
9d ploc-cokkkr (D 118145)
10a
10
Loss
10
Le
10-3
10a
Time (s)
Time (s)
ENGD (Woodbury)
SPRING
Figure 18: Training loss and evaluation 2 error for learning the solution to the 9 1d logarithmic Fokker-Planck equation over time and steps using random search.
Best run details
The runs shown in Figure correspond to the following papeer-prrmmters:
ENGD-W: damping: 1.558 395 x 10-10
SPRING: damping: 7.511 981 x 10-2 momentum: 9.356 251 x 10-1
Search space details The runs shown in Figure 18 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 (l/ denotes uniform, and Ll log-uniform distribution):
ENGD-W: damping: LU((1 x 10-10; 10-1))
SPRING: damping: LU((1 10-10; 10 10-3)); momentum: (Z.90))).999))
30

---

A.6.1 Fixed learning rate
We repeat the previous experiments, now adding search space of [M((1 x 10-!; x 10-4)) for the learning rate, see Figure [9] We find that the best parameters are:
ENGD-W: damping: 8.638 985 X 10-4; learning rate: 6.029 401 X 10
SPRING: damping: 8.377 655 x 10-3; momentum: 9.760 086 x 10-1; learning rate:
4.473 188 10-2
9d lng-fookker-llankk-iitmmpii (/) 118145)
9d Lng-fookkerphannk--oomoii (1) 118145)
10a
10
La error
Loss
10
10
100
10
10a
10a
10
10
10a
10
Iteration
Iteratinn
9d Logg-fokkeerpaamkk-ssttppiii (D 1181455
9d Lage/fokkerplamnk-vooopppii (D 118145)
10a
10a
La emor
Loss
10
10
100
10
10a
10
100
10
10a
10a
Time (8)
Time (8)
ENGD (Wondbury)
SPRING
Figure 19: Training loss and evaluation L error for learning the solution to the 1d logarithmic Fokker-Planck equation over time and steps with fixed learning rate.
B GPU-efficient Implementation Benchmarking
We compare the performance the traditional Nyström approximation [0] and our GPU-efficient proposal. Tests were performed on NVIDIA RTX 6000 GPUs (24 GiB RAM) with Py Porchh's 's built- timing routines, and the per-iteration execution time was averaged over 100 runs with 10 runs as warm-up. We show here scaling analysis for N 5000 and constant regularizer
1.0 x 10-7,
Sketch size 1% of NJ
Avg. Time [s]
Peak Memory [MiB]
Speedup [tmee]]
20%
0.1321 vs. 0.0099
233.9 vs. 165,2
13.29
40%
0.6246 0.0253
395.4 vs. 234.6
24.70
60%
2.1137 vs. 0.0461
599.6 vs. 312.5
45.90
80%
4.7717 0.0725
856.8 vs. 395.0
65.41
Table 1: Nyström approximation performance. Comparison between the standard Nyström
method (10) and our GPU-efficient Nyström (boldfaced values). We report average runtime (s), peak memory (MiB), and speedup (x, defined as Soandard/GUU-efficiinn across sketch sizes (% of N). Our method delivers substantial acceleration (==13-65> -65x and consistently lower peak memory, with both speedups and memory savings increasing with sketch size.
31


Reviews:
[{"review_id": "xayevMO8em", "reviewer": "Reviewer_rnPC", "summary": "The authors use recently developed SPRING algorithm to address the optimization problems in Physics Informed Neural Networks (PINNs). The authors discuss the implementation of SPRING and propose a randomized algorithm to further accelerate the speed. The authors conduct several experiments to demonstrate the effectiveness of SPRING algorithm.", "questions": "If it is possible to implement other kinds of second-order optimizer with woodbury identity?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n1. The authors make a valuable attempt to transfer insights and methodologies from the VMC community to the PINN community, which may boost both fields.\n\n2. The introduction of a randomized variant of the SPRING algorithm offers the potential for improved accuracy-efficiency trade-offs, which may benefit practical applications.\n\nWeaknesses:\n\n1. The experimental evaluation is limited to the Poisson equation, which is relatively simple and may not sufficiently demonstrate the broader effectiveness or advantages of SPRING within the PINN framework. To convincingly establish its utility, the method should be tested on more challenging PDEs, such as the Navier–Stokes equations.\n\n2. The paper directly applies the SPRING algorithm, originally designed for VMC problems, to PINN settings without adequately discussing the fundamental differences and similarities between these domains. For instance, VMC methods draw samples according to a distribution proportional to the square of the neural network, making stochastic reconfiguration (SR)-based methods natural and justified. In contrast, PINNs typically use samples drawn from distributions independent of the neural network, raising questions about why SR-based methods remain effective. A deeper theoretical or empirical discussion on this point is needed to clarify the applicability of SPRING to PINNs.", "was_revised": true, "final_justification": "The reviewer acknowledges that the author's response has adequately addressed the second weakness previously identified. However, the scope of equations studied in the paper remains limited. Therefore, the reviewer recommends a borderline acceptance of this paper.", "rebuttal": "Dear Reviewer rnPC,\n\nThanks a lot for your constructive feedback. Please find our answers below and let us know if you have follow-up questions.\n\n---\n\n> 1. The experimental evaluation is limited to the Poisson equation, which is relatively simple and may not sufficiently demonstrate the broader effectiveness or advantages of SPRING within the PINN framework. To convincingly establish its utility, the method should be tested on more challenging PDEs, such as the Navier–Stokes equations.\n\nThis is a fair point. To address your concern, **we evaluated two more PDEs** (to be included in the next revision): We consider the 4+1d Heat equation and the 9+1d logarithmic Fokker-Planck equation (a challenging non-linear equation). For both equations, we use the same batch size and computational budget as in [1] and the same neural net with 100k parameters. With these settings we tune ENGD-W and SPRING with random search and compare the results to the results of Adam and KFAC from [1].\n\nThe results are as follows:\n\n| Optimizer | $L_2$-error (Heat, 3000s) | $L_2$-error (log-FP, 6000s) |\n| --------- | ------------------ | --------------------- |\n| ENGD-W    | **1.0e-5**        | 1.1e-2                | \n| SPRING    | 1.1e-5             | **1.5e-3**            | \n| KFAC (taken from [1]) | 2.1e-5 | 2.6e-2                | \n| Adam (taken from [1]) | 1.6e-3 | 1.8e-1                | \n\nFrom these additional experiments we conclude:\n1. Both ENGD-W and SPRING achieve state-of-the-art accuracies as in our previous experiments, and the advantage of SPRING is confirmed on the more challenging log-FP equation. \n2. In comparison to KFAC [1], we first note that ENGD-W and SPRING now scale to this setting (before, ENGD in parameter space would run out of memory due to the neural net's size). Furthermore, we find that **both ENGD-W and SPRING reach 2 to 15-times smaller $L_2$-error than KFAC**.\n\nThese new results demonstrate that the applicability of our methods is not restricted to the setting of the Poisson equation. \nLet us know if this makes sense to you and if you have follow-up questions.\n\n---\n\n> 2. The paper directly applies the SPRING algorithm, originally designed for VMC problems, to PINN settings without adequately discussing the fundamental differences and similarities between these domains. For instance, VMC methods draw samples according to a distribution proportional to the square of the neural network, making stochastic reconfiguration (SR)-based methods natural and justified. In contrast, PINNs typically use samples drawn from distributions independent of the neural network, raising questions about why SR-based methods remain effective. A deeper theoretical or empirical discussion on this point is needed to clarify the applicability of SPRING to PINNs.\n\nThank you for raising this point; it is indeed important to clarify why the techniques developed for the VMC community should be applicable to PINNs as well.\n\n**For both PINNs and VMC it is well-established that natural gradient methods are particularly well suited** [2, 3] and frequently outperform stochastic first-order methods like Adam. This motivates transfer between the fields.\n\nOn a more technical level, **the probabilistic structure of the typical loss functions in PINNs and VMC are similar**: in both cases the loss function is defined by sampling from some distribution $p(R)$, which leads to the ENGD/SR update direction being the solution to a least-squares problem involving that same distribution $p(R)$, and as a result these algorithms and their variants such as SPRING should be implemented by sampling from $p(R)$ as well. In the case of VMC the distribution happens to be wavefunction dependent $p(R) = |\\psi(R)|^2$, whereas in PINNs the distribution is a mixture of the uniform distributions over the interior and the boundary of the domain, but this difference does not affect the applicability of SPRING.\n\n\n---\n\n>3. If it is possible to implement other kinds of second-order optimizer with woodbury identity? \n\nYes, in fact Quasi-Newton methods like BFGS use the Woodbury identity by default, see for example the Wikipedia page of BFGS (here called Sherman-Morrison formula, a special case of the Woodbury formula). The vanilla Newton method (using the full Hessian) on the other hand cannot be accelarated with the Woodbury formula.\n\n# References\n\n[1] Dangel, F., Müller, J., and Zeinhofer, M. Kronecker-factored approximate curvature for physics-informed neural networks. Conference on Neural Information Processing Systems, 2024.\n\n[2] Müller, J., and Zeinhofer, M. Achieving High Accuracy with PINNs via Energy Natural Gradient Descent, International Conference on Machine Learning, 2023\n\n[3] Carleo and Troyer, Solving the Quantum Many-Body Problem with Artificial Neural Networks, Science, 2017"}, {"review_id": "Pok3HciVho", "reviewer": "Reviewer_ZTVE", "summary": "The paper presents a number of improvements for training physics-informed neural networks using natural gradient methods.  Precisely:\n\n1. using the Morisson-Woodbury formula to reduce the computational complexity of the matrix inversion from $P^3$ (where $P$ is the number of parameters) to $PN^2$ (where $N$ is the number of vectors in the grid to enforce the DE and boundary conditions).\n\n2. introducing a novel momentum term into borrowed from the SPRING algorithm.\n\n3. further accelerating the matrix inversion using a novel, GPU-efficient Nystrom approximation method to avoid the costly (on a GPU) SVD calculation in the standard Nystrom approximation algorithm.", "questions": "See strengths and weaknesses section.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "I thought that the paper was quite well written overall. It was relatively easy to follow, the contributions were (mostly) clearly described, and the claims made were mostly adequately explained/motivated and experimentally demonstrated. And the experimental results quite clearly demonstrated the benefits in the contribution.\n\nOn the negative side, I think that the second contribution could be more clearly explained. The term $\\phi$ in equation (6) just \"appears\" out of the blue without explanation, leaving the reader to search further afield for clarification. Moreover while I get that (6) is just the usual regularized least-squares expression, what I'm missing is why this is to be preferred over any number of alternative regularization schemes that would, potentially, correspond to different momentum terms?\n\nI am also a tad confused by (9). It appears from algorithm 2 that the $\\lambda {\\bf I}$ term is already incorporated into the GPU-efficient Nystrom calculation, yet for some reason the right-side of (9) reads $(nys (J_k J_k^T) + \\lambda I)^{-1}$ and not $nys (J_k J_k^T + \\lambda I )^{-1}$. Is this a typo, or have I misread the text?\n\nThe results section, however, quite clearly shows the benefit of the algorithm from a computation perspective, so in balance I believe the strengths outweigh the weaknesses here.", "was_revised": true, "final_justification": "The author has addressed my concerns, and I am happy to keep my recommendation as accept.", "rebuttal": "Dear Reviewer ZTVE,\n\nThank you for your constructive feedback. Please find our answers below and let us know if you have follow-up questions.\n\n---\n\n> 1. The term $\\phi$ in equation (6) just \"appears\" out of the blue without explanation, leaving the reader to search further afield for clarification. Moreover while I get that (6) is just the usual regularized least-squares expression, what I'm missing is why this is to be preferred over any number of alternative regularization schemes that would, potentially, correspond to different momentum terms?\n\nIn equation (6), the variable $\\phi$ is just a dummy variable that we optimize over, so it holds no direct meaning. Furthermore, the addition of $\\lambda I$ on the left-hand side of (6) is standard within the natural gradient literature and is equivalent to adding an $L_2$ penalty on the right-hand side.\n\nAside from this, maybe you are referring to the addition of the SPRING momentum term in equation (7)? The intution of equation (7) is to encourage the update direction to be close to the previous one, but it is true that this could be formalized in many ways. We simply use the same choice as the original SPRING paper [1] since this choice has been demonstrated to be effective.\n\n---\n\n> 2. I am also a tad confused by (9). It appears from algorithm 2 that the  term is already incorporated into the GPU-efficient Nystrom calculation, yet for some reason the right-side of (9) reads $(\\text{nys}(J_k J_k^T) + \\lambda I)^{-1}$ and not $\\text{nys}(J_k J_k^T + \\lambda I)^{-1}$. Is this a typo, or have I misread the text?\n\nThanks for catching this. There is a typo on Algorithm 2, which creates the confusion. In the \"Ensure\" statement of Algorithm 2, it should read:\n-  $\\hat{A}_{\\text{nys}} = B B^T$ on the left, and\n- $(\\hat{A}_{\\text{nys}} + \\lambda I)^{-1}v = \\dots$ on the right.\n\nWe will rectify the error in the revised version.\n\n# References \n[1] Goldshlager, G., Abrahamsen, N., and Lin, L. A Kaczmarz-inspired approach to accelerate the optimization of neural network wavefunctions. Journal of Computational Physics, 516: 113351, 2024"}, {"review_id": "RuoaML3MVy", "reviewer": "Reviewer_nNgV", "summary": "This paper introduces techniques for doing second-order optimization for PINN. The method connects previous related methods and theoretical results like ENGD, SPRING, Woodbury's identity, Nystrom approximation, and has demonstrated good performance in terms of convergence speed and accuracy for solving PINNs. Specifically, the paper (1) adapted ENGD and SPRING for the setting of PINN; (2) introduced Woodbury's identity in ENGD to reduce computational and memory cost; (3) proposed a GPU-efficient randomized Nystrom approximation, which is used to randomize ENGD and SPRING. The paper also analyzes the effective dimension that characterizes the limit of randomization. Overall, the paper presents a set of improvements to second-order optimization in PINNs.", "questions": "- Does your method works well with larger networks? All the experiments were conducted the MLP only.\n- For the GPU-efficient randomized Nystrom, do you have any concrete comparsion against the original versino? For example, you could compare wall-clock time and memory usage profiles for different problem sizes, scaling behavior as problem dimensions increase, and other GPU utilization statistics.", "limitations": "No", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n- The method is principled. All the techniques used were mathematically grounded.\n- Main results were good: Woodbury version of ENGD has roughly the same convergence curve but ran faster, and the negative result of randomization is also interesting\n- Experiments were thorough.\n- Technical writing was on point. The presentation of each technical detail was described with sufficient detail without overcrowding the paper.\n\nWeakness:\n- The originality of the paper is a bit lacking. ENGD, SPRING, Woodbury's identity, and Nystrom approximation are all established techniques. The main innovation point of the paper seems to be applying these techniques to the specific setting of PINNs.\n- There are a lot of concepts being shown here. I feel like the organization of the paper could be improved. The abstract could be a bit more descriptive so that the reader has a clearer expectation. The contribution could be clearly stated there.", "was_revised": true, "final_justification": "I raised my score due to the following from the rebuttal\n- The additional experiment indeed demonstrates the efficiency of the GPU-efficient Nystrom approximation\n- The author confirms that the method works beyond MLP", "rebuttal": "Dear Reviewer nNgV,\n\nThank you for your constructive feedback. Please find our answers below and let us know if you have follow-up questions.\n\n---\n\n> 1. The originality of the paper is a bit lacking. ENGD, SPRING, Woodbury's identity, and Nystrom approximation are all established techniques. The main innovation point of the paper seems to be applying these techniques to the specific setting of PINNs.\n\nWe respectfully disagree with the implied concern regarding a lack of originality. According to the NeurIPS reviewer guidelines, originality is not limited to introducing entirely new methods. Rather, *“a work that provides novel insights by evaluating existing methods, or demonstrates improved efficiency, fairness, etc. is also equally valuable.”* Applying Woodbury's identity, SPRING, and the Nystrom approximation to ENGD for the first time aligns with this broader definition.\n\n---\n\n> 2. For the GPU-efficient randomized Nystrom, do you have any concrete comparsion against the original version? For example, you could compare wall-clock time and memory usage profiles for different problem sizes, scaling behavior as problem dimensions increase, and other GPU utilization statistics.\n \nTo corroborate our results, we generate random matrices of size $N=5000$ and compare the standard Nyström implementation of [1] with our GPU-efficient Nyström approximation on a GPU. We run the benchmarking procedure and present the results in the following table (we use the pattern *Standard Nystrom vs. GPU-efficient Nystrom*):\n\n| Sketch size [% of N]    | Avg. Time: Standard vs. GPU-efficient Nyström[s] | Peak Memory [MiB] | Speedup [times] |\n| ---- | ------------------ | --------------- | ----- |\n| 20% | 0.1321 vs. 0.0099  | 233.9 vs. 165.2 | **13.29** |\n| 40% | 0.6246 vs. 0.0253  | 395.4 vs. 234.6 | **24.70** |\n| 60% | 2.1137 vs. 0.0461  | 599.6 vs. 312.5 | **45.90** |\n| 80% | 4.7717 vs. 0.0725  | 856.8 vs. 395.0 | **65.41** |\n \nAs you can see from the results, **the GPU-efficient implementation is at least an order of magnitude faster**. We will include this evidence, and more ablations, in our manuscript.\n\n---\n\n> 3. Does your method works well with larger networks? All the experiments were conducted the MLP only.\n\nYes! Since our methods are linear in $P$ and architecture-agnostic (unlike KFAC), there is no barrier to extending them to networks of larger sizes or different types. The only reason we use MLPs is to keep the networks the same as in other papers on ENGD to enable fair comparisons with previous works.\n\n---\n\n> 4. I feel like the organization of the paper could be improved. The abstract could be a bit more descriptive so that the reader has a clearer expectation. The contribution could be clearly stated there.\n\nWe appreciate the feedback and we would like to make the paper as clear as possible. To help us do so, could you elaborate on what is unclear to you in the current write-up, or what changes you would suggest in order to improve the organization of the paper?\n\n# References\n\n[1] Frangella, Z., Tropp, J. A., and Udell, M. Randomized Nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2):718–752, 2023."}, {"review_id": "VzPQ5lEFwL", "reviewer": "Reviewer_11hA", "summary": "The paper applies a natural gradient optimizer (SPRING) which was originally developed for Variational Monte Carlo (VMC) to PINNs.\nThey find that similar to VMC, natural gradient descent converges substantially faster, and that similar to VMC using the Woodbury matrix identity substantially speeds up the optimizer.\nThe authors additionally try to accelerate the optimizer further by using a randomized Nystrom approximation (with a potentially more GPU-friendly implementation), but find that the randomized approximation is not accurate enough for practical use.\nThe optimizers are tested on the toy problem of the Poisson equation for 5-100D.", "questions": "- What is the performance of KFAC for these systems? In particular how does it fare in the regime of large number of parameters P and large batch size N, where presumable both ENGD (scaling poorly with P) and SPRING (scaling poorly with N) are expensive?\n- Is there a setting where randomized variants of SPRING are superior to deterministic variants, potentially on a more challenging problem which may require larger batch-sizes?\n- Does the faster convergence and runtime of SPRING unlock new applications of PINNs? Can the authors present results on a non-trivial system (e.g. fluid dynamics, inverse problems) where SPRING allows to improve the SOTA?\n\n\nMinor questions:\n\n- What exactly are the \"Naive\" optimizers in Fig. 4,5? Is it SVD-based randomized Nystrom? If yes, why does it converge at basically the same rate as your Cholesky-based Nystrom if the latter one is \"an order of magnitude faster on GPU\"?\n- The authors write on line 189 \"Such a sketch-and-solve approach has significant pitfalls [...]\". What are these pitfalls? Is it \"just\" an inaccurate approximation if the target rank is less than the effective rank?", "limitations": "The limitations section currently only addresses the limitations of the randomized approach.\nI encourage the authors to also address the limitation of having evaluated the optimizers only on very small/simple test systems.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n-------------\n\n- The paper transfers an established \"trick\" to accelerate natural gradient from one field (VMC) to an adjacent field (PINNs). This may be of value to PINNs practitioners, who are unfamiliar with the optimizer literature.\n- On the investigated problems (5D and 100D Poisson equation on a hypercube), SPRING consistently outperforms other optimizers thanks to the convergence rate of natural gradient descent with the speed obtained through the Woodbury formulation\n- The paper is well written and easy to read\n\n\nWeaknesses\n-----------------\n\n- The paper presents minimal novel contributions, and only applies an existing optimizer to PINNs. The novel contribution (randomized Nystrom) is shown to be inferior to the deterministic variant in terms of loss obtained with a given compute budget.\n- The method is only evaluated on extremely simple toy systems (linear PDE with explicit solution)\n- The paper claims that SPRING surpasses KFAC performance, but does not contain any data on KFAC\n\n\nMinor details\n----\n\n- The 2 shades of blue in Fig. 4,9,10 are very hard to distinguish, please consider slightly more distinct colors", "was_revised": true, "final_justification": "I had 2 primary concerns with the paper:\na) Limited evaluation of the method on simple systems\nb) Limited novelty\n\nThe authors have addressed a) by providing data on 2 more systems and also compare against KFAC.\nPoint b) is still a concern for me because Natural Gradient is already widely used in PINNs and simply transferring the Woodbury acceleration trick from VMC seems like a modest contribution.\nIn summary I have increased my rating from 3 to 4 to reflect the authors' improvements.", "rebuttal": "Dear Reviewer 11hA,\n\nThanks a lot for your constructive feedback. Please find our answers below and let us know if you have follow-up questions.\n\n---\n\n> 1. The paper presents minimal novel contributions, and only applies an existing optimizer to PINNs. The novel contribution (randomized Nystrom) is shown to be inferior to the deterministic variant in terms of loss obtained with a given compute budget.\n\nWe respectfully disagree with both A) the implication that randomized Nyström is the only novel contribution of our work and B) the implication that the contribution regarding the randomized Nyström method is not significant due to its underperformance. \n\nRegarding point A), applying the Woodbury matrix identity and the SPRING algorithm within the context of ENGD for PINNS are both novel contributions that yield very encouraging results and have the potential to make an impact on the field of PINNs. Although these contributions do not come in the form of entirely new algorithms, they align with the statement in the NeurIPS reviewer guidelines that **\"originality does not necessarily require introducing an entirely new method. Rather, a work that provides novel insights by evaluating existing methods, or demonstrates improved efficiency, fairness, etc. is also equally valuable.\"**\n\nRegarding point B), the fact that the randomized Nyström approximation does not improve the performance of ENGD is quite surprising and only holds as a result of the algorithmic improvements that we have introduced to the baseline ENGD implementation. In fact, one recent work applied a sketching method to reduce the cost of the *original* ENGD algorithm and found that relative to that inferior baseline, sketching is highly beneficial [2]. Our work thus overturns previous results by showing that, when ENGD is implemented efficiently using the Woodbury matrix identity, randomization is no longer benefcial. \n\n> 2. Is there a setting where randomized variants of SPRING are superior to deterministic variants, potentially on a more challenging problem which may require larger batch-sizes?\n\nWe attempted but failed to identify any scenarios where randomized methods are superior.  As discussed above, this is a surprising finding of our paper which is tightly connected to the algorithmic improvements that we introduced to the deterministic algorithm.\n\n---\n\n> 3. The method is only evaluated on extremely simple toy systems (linear PDE with explicit solution)\n\n\nEchoing our response to Reviewer rnPC, this is a fair point. To address your concern, **we evaluated two more PDEs** (to be included in the next revision): We consider the 4+1d Heat equation and the 9+1d logarithmic Fokker-Planck equation (a challenging nonlinear equation). For both equations, we use the same batch size and computational budget as in [1] sampling one batch at each step, their neural net with 100k parameters, and tune all algorithms with random search.\n\nThe results are as follows:\n\n| Optimizer | $L_2$-error (Heat, 3000s) | $L_2$-error (log-FP, 6000s) | \n| --------- | ------------------ | --------------------- |\n| ENGD-W    | **1.0e-5**        | 1.1e-2                | \n| SPRING    | 1.1e-5             | **1.5e-3**            | \n| KFAC (taken from [1]) | 2.1e-5 | 2.6e-2                | \n| Adam (taken from [1]) | 1.6e-3 | 1.8e-1                | \n\nFrom these additional experiments we conclude:\n1. Both ENGD-W and SPRING achieve state-of-the-art accuracies as in our previous experiments, and the advantage of SPRING is confirmed on the more challenging log-FP equation. \n2. In comparison to KFAC [1], we first note that ENGD-W and SPRING now scale to this setting (before, ENGD in parameter space would run out of memory due to the neural net's size). Furthermore, we find that **both ENGD-W and SPRING reach 2 to 15-times smaller $L_2$-error than KFAC**.\n\nThese findings demonstrate that the applicability of our methods is not restricted to toy examples of linear PDEs with explicit solutions. Let us know if this makes sense to you and if you have follow-up questions.\n\n---\n\n> 4. The paper claims that SPRING surpasses KFAC performance, but does not contain any data on KFAC. What is the performance of KFAC for these systems?\n\nSee the previous response for two new examples where both ENGD and SPRING outperform KFAC. Additionally, for the 100D Poisson equation, KFAC has previously achieved an $L_2$-error of 4.3e-3 with a time budget of 10,000 seconds as shown in [1], while in our work SPRING achieves an $L_2$-error of 3.6e-3 with a smaller budget of only 7,000 seconds and identical hardware.\n\n---\n> 5. In particular how does [KFAC] fare in the regime of large number of parameters P and large batch size N, where presumable both ENGD (scaling poorly with P) and SPRING (scaling poorly with N) are expensive?\n\nYou are right that KFAC scales more favourably in batch size than ENGD-W and SPRING. Our comparison focuses on batch sizes that are used in practise (all settings are from existing papers), and we do not observe this advantage. However, we will make this contrast to KFAC more explicit in the manuscript; thanks for pointing it out.\n\n---\n\n> 6. Does the faster convergence and runtime of SPRING unlock new applications of PINNs? Can the authors present results on a non-trivial system (e.g. fluid dynamics, inverse problems) where SPRING allows to improve the SOTA?\n\nWe view this as an important question for future works; the current work is focused on developing and demonstrating methods rather than tackling specific applications. Nonetheless, we refer the reviewer to the new heat and log-Fokker-Planck equations, which present new, non-trivial settings.\n\n---\n\n> 7. What exactly are the \"Naive\" optimizers in Fig. 4,5? Is it SVD-based randomized Nystrom? If yes, why does it converge at basically the same rate as your Cholesky-based Nystrom if the latter one is \"an order of magnitude faster on GPU\"?\n\nThis is a very good question and we apologize for the confusion. The figures were intended to show the SVD-based method but actually in their current form the naive methods refer to a fully naive Nyström method which approximates $A$ directly as $(A\\Omega)(\\Omega^T A \\Omega)^{-1}(A\\Omega)^T$. This naive implementation is known within the linear algebra community to be unstable, as discussed in [3] and others. In our own (intended to be internal) experiments this method worked surprisingly well (similar to the GPU-efficient stable algorithm), but we do not feel comfortable recommending it due to the known instability, which could become a problem in other settings. Unfortunately, the fix proposed by Frangella et al is not GPU-efficient and performs much worse than both of the methods currently pictured in Figures 4 and 5. This lead us to develop our own Nyström approximation that is both stable **and** GPU-efficient.\n\nRegarding the figures, it was a mistake to include this naive method and we will remove it in the final version. We will instead show the standard (SVD-based) stable algorithm of [3], so that the relative advantage of our proposed GPU-efficient Nyström method becomes clear. We again apologize for the confusion.\n\n---\n\n> 8. The authors write on line 189 \"Such a sketch-and-solve approach has significant pitfalls [...]\". What are these pitfalls? Is it \"just\" an inaccurate approximation if the target rank is less than the effective rank?\n\nYes, the specific issue is that with sketch-and-solve, it is impossible to obtain a highly accurate solution without taking the sketch size to be extremely large, which eliminates any computational savings. We will clarify this point in the manuscript. \n\n\n# References\n\n[1] Dangel, F., Müller, J., and Zeinhofer, M. Kronecker-factored approximate curvature for physics-informed neural networks. Conference on Neural Information Processing Systems, 2024.\n\n[2] Mckay, M. B., Kaur, A., Greif, C., and Wetton, B. Near-optimal sketchy natural gradients for physics-informed neural networks. In Forty-second International Conference on Machine Learning, 2025\n\n[3] Frangella, Z., Tropp, J. A., and Udell, M. Randomized Nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2):718–752, 2023."}]

Author rebuttal:

Meta-review: We thank the authors for their submission.  This work proposes a few practical ways to speed up second order optimization for physics-informed neural networks.  

Reviewers agreed that the paper was well motivated, clearly written, and technically sound, commenting on the claims being well supported by empirical evidence generated.  It is clear that the resulting method has practical benefits. 

Reviewers did voice the concern that this work lacked novelty, as it is a combination of existing methods.  Though this is a valid concern, the impressive performance results and careful empiricism of the work did generate some novel insights.  Additionally, though the approach may combine some well known methods, a degree of technical sophistication appears to be necessary to make them all work together.  Not to mention identifying PINNs as a suitable application to enjoy such benefits.  And the reviewers agreed — though most voiced this novelty concern, all reviewers see the value of the work and ultimately recommend acceptance.
Decision: Accept (poster)