Near-optimal Sketchy Natural Gradients for Physics-Informed Neural Networks
Authors: Maricela Best Mckay
Avleen Kaur
Chen Greif
Brian Wetton
Venue: ICML 2025 poster (2025)
Keywords: Sci-ML
Physics Informed Neural Networks
Natural Gradients
Sketching
Primary area: general_machine_learning->scalable_algorithms

Abstract: Natural gradient methods for PINNs have achieved state-of-the-art performance with errors several orders of magnitude smaller than those achieved by standard optimizers such as ADAM or L-BFGS. However, computing natural gradients for PINNs is prohibitively computationally costly and memory-intensive for all but small neural network architectures. We develop a randomized algorithm for natural gradient descent for PINNs that uses sketching to approximate the natural gradient descent direction. We prove that the change of coordinate Gram matrix used in a natural gradient descent update has rapidly-decaying eigenvalues for a one-layer, one-dimensional neural network and empirically demonstrate that this structure holds for four different example problems. Under this structure, our sketching algorithm is guaranteed to provide a near-optimal low-rank approximation of the Gramian. Our algorithm dramatically speeds up computation time and reduces memory overhead. Additionally, in our experiments, the sketched natural gradient outperforms the original natural gradient in terms of accuracy, often achieving an error that is an order of magnitude smaller. Training time for a network with around 5,000 parameters is reduced from several hours to under two minutes. Training can be practically scaled to large network sizes; we optimize a PINN for a network with over a million parameters within a few minutes, a task for which the full Gram matrix does not fit in memory.

Near-optimal Sketchy Natural Gradients for Plysics-Informed Neural Networks
Maricela Best McKay Avleen Kaur? 2
Chen Greif
Brian Wetton !
Abstract
1. Introduction
Partial differential equations (PDEs) are fundamental in mathematical modelling and simulation. However, most PDEs lack analytical solutions and must be solved numer- ically. Although there are extensive tools for numerically solving variety of PDEs, many problems of scientific inter- est remain computationally intractable or require significant simplification.
Natural gradient methods for PINNs have achieved state-of-the-art performance with errors several orders of magnitude smaller than those achieved by standard optimizers such as ADAM or L-BFGS. However, computing natural gradi- ents for PINNs is prohibitively computationally costly and memory-intensive for all but small neu- ral network architectures. We develop a random- ized algorithm for natural gradient descent for PINNs that uses sketching to approximate the nat- ural gradient descent direction. We prove that the change of coordinate Gram matrix used in a natu- ral gradient descent update has aapidly-decaying cigenvalues for a one-layer, one-dimensional neu- network and empirically demonstrate that this structure holds for four differrnnt example prob- sketching algo-
The success of deep learning has generated great interest in combining deep learning techniques with scientific com- puting (Weinan et al., 2021; Baker et al., 2019). Physics- informed neural networks (PINNs) (Raissi et al., 2019) are one such technique. PINNs have appealing features for tackling challenges in scientific computing and modelling, and have shown a lot of promise across broad range of applications; they have been used for data assimilation, to model multi-scale and multi- multi-physics phenomena, inverse problem solvers for parameter estimation, solve high- dimensional systems of PDEs, and to combine incomplete mechanistic understanding with data (Karniadakis et al., 2021; Dabrowski et al., 2023).
lems. Under this structure, rithm is guaranteed to provide near-optimal low- rank approximation of the Gramian. Our algo- rithm dramatically speeds up computation time and reduces memory overhead. Additionally, in our experiments, the sketched natural gradient outperforms the original natural gradient terms of accuracy, often achieving an error that is an order of magnitude smaller. Training time for network with around 5,000 parameters is reduced from several hours to under two minutes. Training can be practically scaled to large network sizes; we optimize a PINN for a network with over a million parameters within few minutes, a task for which the full Gram matrix does not fit in memory.
Despite their demonstrated utility, important limitations have emerged. PINNs can be notoriously difficult op- timize, particularly for multi-scale problems and problems that exhibit multiple frequencies (Rohrhofer et al., 2022; Wong et al., 2022; Krishnapriyan et al., 2021; Wang et al., 2021; 2022). First-order methods such as ADAM (Kingma & Ba, 2015), while scalable, are generally unable to achieve appropriate performance accuracy. While memory-intensive second-order optimizers such as BFGS or L-BFGS often fare better (Krishnapriyan et al., 2021; Zhu et al., 1997; Rathore et al., 2024), PINNs rarely achieve accuracy better than 1E-03 (Müller & Zeinhofer, 2023).
Recently, another class of optimization methods, natural energy gradients, has been developed for PINNs to address these difficulties (Müller Zeinhofer, 2023). Natural en- ergy gradients take advantage of the underlying geometry in the function space for the PDE residual related to PINN to formulate a better search direction in the parameter space. These methods show significant promise; they can achieve state-of-the-art accuracy, often outperforming other opti-
'Deaarmmnt of Mathematics, University of British Columbia, BC, Canada *Department of Computer Science, University of British Columbia, BC, Canada. Correspondence to: Maricela Best Mckay <maiccola @ math.ubc.ca>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Near-optimal Sketchy Natural Gradients for PINNS
mizers by orders of magnitude. However, natural gradients prohibitively computationally costly for all but small network sizes.
Wu al. (2023); Adcock et al. (2022); Daw al. (2023) study strategies for sampling the collocation points used in training to improve training.
In this paper, we aim to scale these methods while main- van der Meer et al. (2022) focuses on adaptively weigh- ing different components of the PINN loss.
taining the high accuracy that they provide. We introduce sketchy natural energy gradient descent (SNGD), a random- ized algorithm for natural gradient descent for PINNs that uses sketching (Tropp al., 2017) approximate the natu- ral gradient descent direction.
Wang et al. (2024) forces a PINN to respect causality by modifying the PINN loss to enforce temporal order in time-dependent problems.
Wang et al. (2021) suggests specialized architecture to improve training.
Main contributions Our main contributions can be sum- marized as follows:
Various other optimization strategies have been sug- gested for PINN training (Zeng et al., 2022; Davi & Braga-Neto, 2022; Nurbekyan et al., 2023). These strategies have worse performance accuracy compared to both ENGD and the present work.
Our algorithm is fast, scalable and accurate. We demon- strate that method dramatically speeds up compu- tation time and reduces memory overhead by training a network with over a million parameters within a few minutes. For this network size, the full Gram matrix required for natural energy gradients cannot fit in mem- ory. In our experiments, SNGD consistently achieves error rates that are more than an order of magnitude lower than those of the original natural gradient.
Rathore et al. (2024) provides an in-depth analysis of optimization challenges in PINNs. They examine the ill-con ill-conditioning of the loss landscape, which is linked to the spectral properties of the Hessian, and intro- duce a novel second-order method for PINN training, NysNewton-CG.
Our method exploits structural properties inherent to the change of coordinate Gram matrix used in natu- ral gradient descent update. In Section 3, we provide insight into the structural properties of the change of coordinate Gram matrix used in natural gradient de- scent update. We provide empirical evidence that this matrix is numerically low-rank, ill-conditioned and has rapidly-dccaying cigenvalues. We show that for a one-layer, nie-dimeniional neural network, the eigen- values of the Gramian decay rapidly. In Section 4. we show that under this structure, our sketching algo- rithm is guaranteed to provide near-optimal low-rank approximation of the Gramian.
Another line of related work is the use of sketching for higher-order optimization.
Yang et al. (2022) introduces sketching techniques for approximating the empirical Fisher information matrix (EFIO). However, the problem under consideration, as well as the use of sketching to accelerate natural gradient descent, differ significantly from the treatment presented in this work.
Gower al. (2019) proposes randomized subspace Newton, which uses sketching to approximate the Hes- sian of the loss function. The authors focus on linear regression problems using sketching to approximate the Hessian. This differs from the present work, which considers non-linear optimization problems and natural gradient methods rather than Newton's method.
Our method is a computationally efficient technique for promoting the training of higher-frequene, com- ponents in neural networks and thus is relevant to the problem of spectral bias exhibited by fully connected neural networks (Rahaman et al., 2019; Zhang et al., 2023). See Appendix A and Appendix B for more insight.
2. Background
We are interested in approximating, with a neural network, the solution to the following system of differential equa- tions,
Related work Many works address the accuracy and train- ing difficulties in PINNs. Some of which are complementary to the present work in that they can be used together with SNGD.
Du = f in Ω
Bu g on an
(1)
Müller & Zeinhofer (2023) propose energy natural gra- dient descent (ENGD), which the current work directly builds on.
where D(-) is a differential operator, B(-) is a boundary value operator, Ω C RR, aQ is an appropriately defindd
2

---

Near-optimal Sketchy Natural Gradients for PINNS
boundary, and u is the unknown solution. I and g are known functions in I22. Note that 9 E 82 may include initial condition for time-dependent system. Consider the form
where |0] is the number of parameters in the neural net- work, and (Ōo, D) denotes the composition of the partial derivative of ug with respect to its i-th component with the differential operator D.
F(u, v) = / (Du-f)(2 - f) dΩ
3. The structure of the Gram matrix
/ (Bu (B 9) ds.! (2)
The Gram matrix G is a dense, symmetric, and positive semi-definite. It can also be observed empirically that G is numerically low-rank, ill-conditioned, and has rapidly- decaying eigenvalues. This empirical observation is con- sistent across all the experimental results presented in this paper. Figure 1 shows an illustrative example of this be- havior for a neural network with 1,341 parameters taken from run 2 of the experiments described Section 5.4. Eigenvalues are shown for every 400 iterations of training.
Finding the solution u to Equation (1) is equivalent to find- ing R(u)= = F(u, u) = O.
In PINNs, we aim to approximate the solution u to Equa- tion (I) with some neural network, which we will call ug. by minimizing
R(uo)= -/1////0 (Duo- - f)² dit+ S (Buo - 9)² ds, (3)
Low-rank The success of deep learning has gone hand in hand with the over parameterizationn of neural networks (Sejnowski, 2020; Belkin et al., 2019; Chang et al., 2021). With this in mind, it is advisable to choose a network ar- chitecture that ensures over peee-parameterzzaionn for the target task; this design paradigm implies that the matrix G should be rank- deficiefn..
where 0 is the set of parameters of the neural network. The derivatives of Wg can be easily computed with automatic differentiation. The norms in Equation (3) are replaced by Euclidean norms with an appropriately-defiee quadrature, which defines loss function on the parameter space of the network. To find numerical approximation to Equation (1), the parameters 0 are trained using optimization algorithm, such as gradient descent.
Figure/ demonstrates the rank deficiency of G, for the experiment depicted |0| 1, 341 but G has only around 800 eigenvalues above machine precision at iteration 2,000 and can be observed to be rank deficient throughout training.
In natural energy gradient descent (ENGD), the traditional gradient typically used in a gradient descent step is replaced with a new gradient direction that is related the function space geometry of the loss (Müller & Zeinhofer, 2023). Let R denote the discretization of R using an appropriate quadrature rule, then the loss function is
The Gram matrix is ill-conditioned be- cause there are small non-zero eigenvalues close to machine round-off. We can observe in Figure that a significant portion of the eigenvalues of G become very small and so G is severely ill-conditioned. See Appendix A and Ap- pendix B as well as Zhang et al. (2023) for insight into how this phenomenon is related to the spectral bias of neural networks.
Ill-oondtiiondd
I(0)= = R(us).
A natural energy gradient descent update is given by
0++1 = 02 - G/ L()).
(4)
where 1) is the step size, denotes the Moore-Penrose pseu- doinverse, h E N is some iteration of the descent, and
Rapid eigenvalue decay In all of the experiments in this paper, G exhibits exponential eigenvalue decay throughout the training process. At each iteration, the eigenvalues de- cay machine unit roundoff an exponential rate. This behavior can be clearly seen in Figure 1.
v L= = (0)Vo
(5)
is called the natural energy gradient?.
Entries of the Gram matrix G are given by an appropriate discretization (via quadrature) of
For one-onnensseeeeee dieedimensional, one-layer neural network with ReLU activations, following the construction in Zhang et al. (2023), the j-th largest eigenvalue of the Gram kernel function for the network decays approximately as j i.e., Aラ ~ j-4, This result can be seamlessly extended to show exponential decay for the Gram kernel function in one dimension (in continuous form), which corresponds to (6) in its semi-
G ( (10e. o ())b (uol) dx
In
+ I ((ob o on
B) (uol) ds,
$ و < 10), (6)
In the case D is linear and f = g 0, Fuu, v) is inner- product.
Notice that replacing G with the identity in Equation (4) results a normal gradient descent step.
3

---

Near-optimal Sketchy Natural Gradients for PINNS
discrete form.
The key idea is replace G with a much smaller "sketched" matrix. Concretely, G is replaced with A GM where M is rectangular matrix of size |0] x L, where < 10), whose entries are sampled from a Gaussian distribution.
Lemma 3.1. For the Gram kernel function G, define the operator K= ([a bl) - (2 b))
Kh(a)- = =f" (" w)l
To compute a sketched natural gradient update,
Ok+1 = 02 - V L,
denote >j as the j-th largest eigenvalue of K. Then the eigenvalues } decay exponentially.
we employ single-pass and two-pass sketching techniques from Saibaba et al. (2016). Both find the randomized sin- gular value decomposition (RSVD) of a matrix. Due to the symmetry in our case, this amounts to a randomized eigenvalue decomposition. We regularize by truncating the cigenvalues below specified tolerance to robustly approxi- mate the pseudoinverse of G (see Algorithm 1). A unique feature of our setting is that G possesses exponentially- decaying eigenvalues, which guarantees that the error can be driven close to machine precision for single/two-pass RSVD (see Section 4.1).
For any analytic activation function, 9 is an analytic fune- tion. Lemma 3.1 is a direct implication of the well-known result that the eigenvalues of an integral operator with ana- lytic kernels decay exponentially (König & Richter, 1984, 148). This proves our claim as cigenvalues of the Gramian matrix G and the corresponding integral kernel differ by a constant (Zhang et al., 2023, p. 9). For more details, see Appendix A and Appendix B.
10
Amma
(1/181A
Instead of generating the entire Gram matrix and then mul- tiplying it by M, it is possible to proceed more efficiently. For any matrix Y we can compute GY as the discretization of
10 18
Eigervalues [Ax] (log scale)
10"-
0g
i=D
10"
i= 400
800
10-3
1200
Gi = ((00). o D) (Y" DD]uoll, d 22
1600
10-9.
2000
10-9
/ ((0b, B)(o))) (Vo B) [uol] ds, 80
10-12
1 s i S 10), (7)
10-1b
under an appropriate quadrature. This approach has the advantage of requiring only the computation and storage of the [0) x entries of A, as opposed to the computation and storage of the 10]2 entries of G in addition to the matrix- matrix multiplication GM.
0
100
600
800
Index of eigenvalue Ae
Figure 1. Eigenvalues' of the matrix G at iteration of training vs. the index of the eigenvalue shown with the line y = Amaxe + This figure is from run 2 in Section 5.4(sce Table 6). Figenvalues are considered zero at 1E-16 and below this threshold. Note that the number of non-zero eigenvalues changes during training. A horizontal line at 1E-13 is also shown; this corresponds to the tol value used for sketching throughout this paper (see Algorithm for details).
Computational complexity and memory requirements We briefly discuss the per-iteration differences in com- putational complexity and memory requirements between ENGD and SNGD. The memory requirements for ENGD are (10 whereas for SNGD they are ((1 p)), where (r + p) << [0]. The computational cost for com- puting the matrix G is variable and depends on the structure of the underlying PDE operator. For the full matrix G, the cost is a function of the number of nested autodifferenti- ations, which must be computed for each data point and distinct entry in G. Therefore, using the sketched matrix via Equation (7) leads to significant savings of factor of ((00/(r + p)) for this part of the computation.
4. Sketching process
To compute V L defined in Equation (5), it is necessary to compute G as well as the action of its pseudoinverse. Since G is a square matrix with 2 entries, both of these operations can be quite computationally costly and memory- intensive for networks of even moderate size. To address these challenges, we employ randomized numerical linear algebra techniques (Halko et al., 2011).
The remaining computational complexity for each algo- rithm is related to computing GiV In ENGD this is accomplished with a least-squares solve at a cost of ()))3). SNGD, on the other hand, has a computational cost of
Onnly every 15th eigenvalue in Figure is shown in order to aid in plot legibility.

---

Near-optimal Sketchy Natural Gradients for PINNS
Algorithm 1 Sketchy Natural Gradient Descent (SNGD)
We know that the eigenvalues of G decay exponentially, so
Input: Initial parameters 0, initial rank estimate r, over- sampling parameter p. truncation threshold tol, and max- imum iterations Nmax- for i to Naxx do Compute (0,(0;)
X3
2 A
j=r+1
Sample M E I ((x((tr) from a Gaussian distribution via Equation (7)
Therefore, (8) becomes
A <- GM
QR - A
Reduced QR factorization
E||G G C(p,r, & 10-1
(9)
Q - Q(:, :: rr)
Compute T:
where C(p, r,e) is a positive constant that depends on p, 7, and
Single-pass: T Q"
Two-pass: T + GQ
via Equation (7)
Note that although the constant C(p, e, is polynomial in r, for sufficiently large the eigenvalues decay exponentially, i.e rapidly enough to beat the polynomial growth in C. A minimal amount of oversampling drives the error estimate to the theoretically minimal value, and we obtain a nearly optimal estimate G R G (ptrr = (QS)²' (QQS), with significantly fewer computation and memory requirements as compared traditional truncated SVD techniques applied to the full matrix G.
SA SArtp) ST =T
+ [Mip, > tol
D + diag((1/Auss
U QS
A L(,,)) UD'
1/" f arg 10 L (0, -0:-n L(O.))
0, = 0 -
end for
Why not use CG? Conjugate gradient (CG) can be used for solving least-squares problems; however, for dense posi- tive semi-definite low-rank matrices with rapidly decaying eigenvalues such G in the present work, CG is not an ideal choice because it is based on matrix-vector products and the presence of near-zero eigenvalues slows down its convergence.
(10((p + r)) (see Algorithm 1). Both the computational complexity and memory requirements are dynamic and change the estimated rank r changes.
4.1. Near-optimal error when sketching G
As discussed in Section 3, G is a symmetric positive semi-definite matrix, which is dense, low-rank, and ill- conditioned with rapidly decaying eigenvalues. We have used a randomized eigenvalue decomposition with single- pass* and two-pass techniques to estimate the sketched natural gradient direction S which depends on the pseudoinverse of G (Saibaba al. (2016, p. 316)), de- noted by G The techniques mentioned above give us G G G (ptr) == QSA(QS)) where Q and S have dimen- sions 10] x (p + r) and (p rr x (p + r), respectively, and possess orthonormal columns. Halko et al. (2011, p. 273) de- rive the following average spectral error for approximating G by G ((ptr).
Why does sketchy natural gradient descent outper- form natural energy gradient descent? Due to the ill- conditioning of G. we seek a numerically stable algorithm. Randomized algorithms putationally efficient and cut off near-zero cigenvalues. We further truncate any eigenvalues that are below specified tolerance. Comparing the rank of G at the end of training under SNGD vs. ENGD, computed as the number of eigen- values above our tolerance, we observe that under SNGD, G has larger final rank (see Section 5.2 and Section 5.4). This suggests that SNGD finds better and flatter local minimum, which is phenomenon that has been observed to be useful for generalization and related to spectral bias (Li et al., 2018; Fridovich-Keil et al., 2022). We believe that the randomness introduced by multiplying by a distinct Gaussian matrix at each training iteration, as well as the row-mixing effect of this multiplication, contributes the enhanced performance (see, e.g., (Avron et al., 2010) for illuminating discussion of the effect of mixing the coherence of a matrix).
suitable because they com-
/ r
B|| G G(ptr)ll < (1+
Ar41
V p-1
eVp+r
12
(8)
P
Note that (8) is mentioned for randomized SVD, which can be extended randomized eigenvalue decomposition by Section 5.3 the reference.
"also in Halko et al. (2011, 251).
5

---

Near-optimal Sketchy Natural Gradients for PINNS
5. Experiments
5.1. The heat equation
We contrast the results of training a PINN using the original natural energy gradient descent with our sketchy version on several problems. We first reproduce the experiments from Müller & Zeinhofer (2023) and compare them with our method, and then include new problem: the transport equa- tion with high-wave speed. This problem is deceptively sim- ple yet notoriously challenging for PINNs (Krishnapriyan et al., 2021).
To begin, we consider a simple example: the one- dimensional heat equation with Dirichlet boundary con- ditions and simple source term given by
Ou 10uu
for (f, E Ω [0,1] x [0.1]
Ot 0x²
m(0,, x) sin (Tx)
for x E [0. 1]
u(t, x:) 0
for x), E [0, (0. 17. (10)
The form for this problem is given by:
Flu v) = J/ u( (000x)0000, dx +
f
Implementation details All of the code is implemented in Python using the Jax library (Bradbury et al., 2018). The neural networks are built using the Equinox library (Kidger& & Garcia, 2021). The Optax and Jaxopt libraries used for ADAM and BFGS, respectively (Babuschkin et al., 2020; Kingma& & Ba, 2015; Fletcher, 1987). All experiments were run with Google Colaboratory using an NVIDIA A100 GPU. All experiments were run using double precision. We have found that this is necessary to ensure adequate accuracy in computing G. Code to reproduce experiments in this manuscript is available at http://gthhbboommMMaaclla//////
we ds dl.
Ixn
+ -/ 1/2
(u
10au
(v
1 100u
dx dt.
00
4
xxx
01
4 dxx²
The related natural energy Gram matrix is
Gi,- = F Ouo Ouo a0, 00,
where we again use F to denote the discretization of F, where f related to the form F in that is the same object with the linear terms f and g dropped. This is due to the linearization w.r.t. the parameters that occurs in Equation (6).
Choosing the hyperparameters in SNGD Algorithm requires selecting three hyperparameters: an estimate for the rank of G at initialization, the (constant) oversampling parameter p. and a tolerance for the truncation threshold for the randomized eigenvalue decomposition of G.
We train three different PINN architectures to approximate the solution u((,a:) = e (xxxt/4) sin अह to Equation (10). The first network is a one-layer network with 261 parame- ters, the second is a three-layer network with 5,301 parame- ters, and the third is a five-layer network with over a million parameters (referred as A1, A2, and A3, respectively in Table 1 and Table 2). The first architecture is selected be- cause it is similar to the one-layer 257 77-paaameter network used in Müller & Zeinhofer (2023). The other two archi- tectures are selected to showcase how SNGD scales as the number of parameters grows.
To initialize r for Algorithm 1, a single/two-pant step is performed with hand-selected sketch size. The initial rank estimate is chosen to be the number of eigenvalues above specified tolerance for this initial sketched matrix. Rather than hand-tunin -uning the constant p. we use the aforementioned hand-selected sketch size and set p equal to the first index for which the cigenvalues of the sketched matrix are below machine precision. As training progresses, the rank of G grows, reaching maximum value before plateauing to stable value the loss reaches the local minima (this behavior can be observed in Figure 1).
As reported in Table 1, SNGD achieves better accuracy than ENGD, often achieving over a magnitude or, in some cases, nearly two orders of magnitude smaller error. Compared to training with ADAM, SNGD is able to improve accuracy by around three orders of magnitude. In Figure 2 we see that SNGD is able to reach a better minima than the other optimizers, including ENGD, Table shows that networks of around five thousand parameters reach similarly accu- rate minima compared to networks of around one million parameters for PINNs approximating the heat equation; we include results for both networks to showcase the scalability of SNGD to large parameter spaces.
In practice, we have observed that Algorithm 1 is quite sensitive to the choice of tolerance. Because G scales the gradient of the loss, this parameter captures the (locally) important directions in parameter space. There is tradeoff between the need to optimize speed of convergence, which would require tol to be sufficiently large, and the need to ensure accuracy as stated in Equation (9), which requires tol to be small. Equation (9) and the exponential decay of the eigenvalues ensure that the bound on the error will be smaller than tol. We thus choose tol to be 1E-13, which we have found to be performant and accurate enough.
'G can be expressed in this simpler form because the heat equation a linear PDE.
6

---

Near-optimal Sketchy Natural Gradients for PINNS
Table 1. Relative L2 emrors for training a PINN to approximate the heat equation using different network architectures, averaged over five random initializations. The first two architectures are trained for 2,500 iterations for ENGD and SNGD, and 3,000 for BFGS. The last architecture is trained using SNGD for 1,000 iterations. For this last example, ENGD and BFGS are infeasible. ADAM is run for 100,000 iterations for all network sizes.
Table 2. Total computation times (in minutes) for training the heat equation using the three PINN architectures averaged over five random initializations. Training time for ENGD and BFGS for the third architecture with over million parameters is omitted because is not possible to train this PINN with these methods as the memory requirements surpass system availability.
OPTIMIZER
A1
A2
A3
OPTIMIZER
MEDIAN
MIN
MAX
ADAM
24.87
44.15
64.48
A1: 261-PARAMETERS
BFGS
1.10
1.86
N/A
ENGD
1.34
221.89
N/A
ADAM
1.39E-03
8.69E-04 04
1.65E-03
SINGLE-PASS
1.06
1.56
16.663
BFGS
1.99E-04
5.0E-05
1.77E-01
TWO-PASS
1.10
1.86
11.283
ENGD
2.78E-06
1.95r-06
8. 8.94r-06
SINGLE-PASS
3.62F-07
8.23F-08
5.62F-06
TWO-PASS
1.49F-07
1.27P-07
4.86F-06
5.2. Poisson's 's equation
A2: 5,301-PARAMEES PARAMTTERS
The next example the 2D Poisson equation,
ADAM
6.13E-05
2.59E-05
1. 1.
BFGS
3.65E-06
2.32E-06
6.50E-06
Vu(a, y) f(z. y)
for (a, y) = [0, 1]2
ENGD
5.66E-07
2.91E-07
1.36E-06
(11)
u(x,y) 0
for (x, 0Ω,
SINGLE-PASS
7.87E-09
6.00E-09
1.68E-08
TWO-PASS
1.24E-08
8.7E-09
1.6988-08
where, f(x, y) 2n sin(xr) sin(xy) and the solution is u(x, y) sin (T.r) sin(Ty))
A3: 1,085761
PARAMETERS
ADAM
1.10E-05
8.42F-06
2.70F-05
SINGLE-PASS, PASS
1.484E-08
8.92E-09
2.65E-08
10
TWO-PASS
1.49E-07
8.63E-09
1.99E-08
10-1
10-7
10
10-
10-11
Loss
10
500
25003
30000
75000
Iterations
Iterations
10-9
10-12
BFGS
Two-pass ENGD
ADAM
1000
50000 7X00 100000
Single-pass
Iterations
Iterations
BFGS
Two-pass ENGD
ADAM
Figure 3. The loss averaged across five random initializations vs. iterations for the poisson equation, shown for the five different optimizers.
Single-pass
Figure The loss averaged across five random initializations vs. iterations for architecture A2 for the heat equation, shown for the five different optimizers.
We train a PINN with four hidden layers and 1, parame- ters approximate the solution Equation (11). As reported in Table and Table 4, SNGD reaches accuracy over an order of magnitude better than ENGD in significantly less computation time. We also observe from Figure 3 that SNGD reaches a different minima the loss landscape from ENGD. The loss is smaller under SNGD by the end of training, and the number of parameters that meaningfully contribute optimization is greater for SNGD. The trun- cated rank is reported in Table 4, this value is measure of the eigenvalues of G that are above tol = 1E-13. This indicates that the minima for SNGD has a different local curvature than that of ENGD and that more of the dimen- sions of the parameter space contribute to the shape of the local curvature.
In Table 2, we report the total computation time averaged across five different random initializations for the experi- ments Table 1. SNGD comparable in computational cost to ENGD for very small network sizes but is able to optimize more realistically sized networks within minutes of computation time. The network with around 5,000 pa- rameters takes under two minutes to train using SNGD compared to nearly four hours for training the same network via ENGD. For the network with over one million parame- ters, SNGD takes less than a fourth of the time of ADAM to achieve an error that three orders of magnitude better.

---

Near-optimal Sketchy Natural Gradients for PINNS
rameters approximate the solution w(x) cos( to Equation (12). We report our results in Table 5. ENGD and SNGD have similar accuracy for this problem, both order of magnitude more accurate than BFGS. However, SNGD is more accurate than ENGD and requires around one fourth of the computation time for this network architecture.
Table 3. Relative L errors for training PINN to approximate the Poisson's 's equation, averaged over five random initializations. ENGD and SNGD are trained for 2,000 iterations, while BFGS is trained for 3,000.
OPTIMIZER
MEDIAN
MIN
MAX
ADAM
4.67E-04
2.89E-04
6.95E-04
5.4. The transport equation
BFGS
8.0E-06
4.92E-06
1.05E-05
ENGD
5.16E-07
3.78E-07
7.27E-07
We turn our attention to the one-dimensional transport equa- tion with periodic boundary conditions and a sine initial condition.
SINGLE-PASS
9.46E-09
7.10F-09
1.96E-08
TWO-PASS
1.21E-08
7.58E-09
1.98-F08
Ou
du
Table 4. Loss, total computation time (in minutes), and truncated rank of the matrix G for training a PINN approximate Poisson's equation averaged over five random initializations. ENGD are trained for 2,000 iterations, while BFGS is given 3,000. The rank is truncated at fol = IE-13 and computed at the end of training.
+ c
=
for (z, ().t) E
Ōt
dr
(13)
w(0, t) = u(2T, t) for x E 0n
" (x 0) = sin(x)
where, Ω = [0, 2n] x [0, 1] and an = {0, 2n). This ex- ample models straightforward phenomena; it transports (or convects) the initial condition with speed in the direction of the sign of The analytical solution to this problem u(I, t) = sin(2 ct) can be casily found. However, despite the straightforward nature of Equation (13), for values of c >> 10, PINNs fail to converge (Krishnapriyan et al., 2021; Daw et al., 2023).
OPTIMIZER
Loss
RUNTIME (MIN)
RANK
ADAM
2.388-05
52.51
N/A
BFGS
1.46E-8
0.96
N/A
ENGD
1.55E-10
8.85
105.6
SINGLE-PASS
1.84--14
1.52
160.6
TWO-PASS
5.23E-14
1.92
134.6
5.3. A non-linear boundary-value problem
We choose wave-speed 30 and train PINN with four hidden layers and total of 1,341 parameters using two- pass SNGD, ENGD, and BFGS. In Table 6, we present the relative L error and the value of the loss achieved by the three optimization methods. For SNGD and ENGD, we also show the rank of the matrix up to tol = 1E-13.
We now consider a non-linear problem:
[- u" + 3² = x² cos(TTz) cos (πι) for x E === =0==0 u' (x)) =
for x. E 1-1,
(12)
In Müller & Zeinhofer (2023), small neural networks with at most 257 parameters are used train ENGD. Using such small network is insufficient for capturing the essential features of this example. We have observed that both SNGD and ENGD fail to converge during training for compara- ble networks. Additionally, in Table 6, we can see that by the end of training, G has more than 257 significant eigenvalues. These eigenvalues correspond directions in the loss landscape that are scaled by the Gramian matrix (see Appendix A for insight into how G affects the descent direction).
For this example, variational formulation of the problem is used to calculate the form F from Equation (2).
Table 5. Relative L errors and total computation time for training a PINN to approximate the non- non-linear boundary value problem averaged over five random initializations. ENGD and SNGD are trained for 1,000 iterations, BFGS is trained for 3,000, and ADAM given 100,000 iterations.
OPTIMIZER
MEDIAN
MIN
MAX
As reported in Table 6, ENGD fails to converge for runs 1, 3, and 4. In run 1, the loss reduced without making progress in reducing error, suggesting that the network converging to a trivial solution or fixed point for the differential operator in Equation (13) (Rohrhofer al., 2022; Wong et al., 2022). When ENGD succeeds at training, i.e. for runs 2 and 5, it achieves comparable accuracy to BFGS. These errors are similar to those achieved in Krishnapriyan et al. (2021) and Daw et al. (2023) (curriculum regularization and R3 sampling). The loss for ENGD, when it converges, is also similar to BFGS, as can be seen in Figure 4. SNGD, on the
ADAM
2.288-05
5.32E-06
4.64E-05
BFGS
1.20E-07
3.52E-08
2.43E-07
ENGD
7.39E-08
1.28E-08
2.31E-07
SINGLE-PASS
3.66F-08
1.34F-08
5.56F-08
TWO-PASS
4.02r-08
1.41F-08
6.28F-08
RUNTIME (MIN)
BFGS
ADAM
ENGD
SINGLE-PASS
TWO-PASS
1.88
52.51
4.56
1.09
1.32
We train a PINN with four hidden layers and 1,341 pa-
8

---

Near-optimal Sketchy Natural Gradients for PINNS
102
Table Comparison of training PINN for the transport equation using SNGD and ENGD for five different random initializations. Relative L errors, the value the attains, and the rank the matrix G. up cut-off tolerance of IE-13 the end of training are shown. methods are trained for 3,000 iterations.
10
88
10
10-
RUN
ERROR
Loss
RANK
2000
S0000
3000000
Iterations
Iterations
SINGLE-PASS SNGD
BFGS
Two-par ENGD
ENGD (converged)
Single-pass
ADAM
1
4.59E-04
1.61E-11
661
2
9.51E-04
1.10F-13
484
3
1.23E-03
1.52E-13
551
Figure 4. The averaged across five random initializations vs. iterations for the transport equation, shown for the five different optimizers. The average across five random initializations for ENGD is shown in red. Because the loss did not make any progress for two of the random initializations, the average across the three random initializations for which the loss converged is also shown in purple.
2.01E-03
9.98E-14
$20
5
2.45E-03
7.68E-13
533
TWO-PASS SNGD
1
2.99E-4
2.94E-9
2
7.24F-4
2.70F-11
3
3.78E-4
9.53E-10
341
4 8.13F-4
2.51F-11
257
5
2.73E-4
4.90E-12
299
ENGD
1 3.33E0
4.68E-7
246
2
2.85E-3
3.36E-8
187
3
1.82BO
7.37E0
89
Instead of explicitly forming the Gram matrix G and per- forming a matrix-matrix multiplication with M, we com- pute directly by using (7), which only requires storing [0] x entries rather than the full |0]2 matrix. This memory- efficient implicit computation of SNGD has enabled us to train PINNs with millions of parameters, which is otherwise infeasible for ENGD, which requires storing and computing the full matrix G. Our method achieves this within a few minutes, stark contrast to the several hours of runtime required by existing approaches. Furthermore, our error considerably lower than that of the current state-of-the optimization schemes for PINNs.
4 1.79H0
5.03EO
83
5
2.33E-3
2.58E-8
167
BFGS
3.32E-3
2.65E-7
N/A
2 7.68E-2
1,14E-10
N/A
3 2.24E-3
1,41E-7
N/A
4 3.67F-2
1.34R-9
N/A
5
7.67F-3
5.29R-9
N/A
ADAM
1 1,29E+00
2.858-02
N/A
2
1.37E+00
3.03E-02
N/A
1,24E+00
2.68E-02
N/A
4
1.35E+00
3.20E-02
N/A
5
1.27E+00
2.77E-02
N/A
Impact Statement
This paper presents work aimed at advancing the field of Machine Learning. There are many potential societal conse- quences of our work. We have developed computationally efficient technique to promote the training of oscillatory components (high wave number/frrequency) in neural net- works. Additionally, algorithms that enhance the accuracy and computational efficiency of PINNs have numerous po- tential consequences due to the myriad applications of sci- entific interest.
other hand, trains to an order of magnitude better accuracy (or more) across all five random initializations, which, to the best of our knowledge, is state-of-the-an for this benchmark problem.
6. Conclusion
We have designed SNGD, tailored randomized algorithm that exploits the structure the Gramian matrix G. which positive, sefi-iefinite numerically low rank, and ex- hibits exponentially decaying eigenvalues. The error bound (8) suggests that randomized cigenvalue decomposition is highly suitable for approximating G: both the single- and two-pass methods significantly outperform the tradi- tional least- cautre schemes used in ENGD.
A Acknowledgments
We thank the reviewers for their helpful suggestions, which improved the clarity of this work. We gratefully acknowl- edge financial support from the Natural Sciences and Engi- necring Research Council Canada (NSERC). CG and AK were supported by RGPIN-2023-05244,

---

Near-optimal Sketchy Natural Gradients for PINNS
References
Dabrowski, J. J., Pagendam, D. E., Hilton, J., Sanderson, C., MacKinlay, D., Huston, C., Bolt, A., and Kuhnert, P. Bayesian physies informed neural networks for data assimilation and spatio-tempore modelling of wildfires. Spatial Statistics, 55:100746, 2023.
Adams, R. Sobolev Spaces. Academic Press, Inc., 1975.
Adcock, B., Brugiapaglia, S., Dexter, N., and Morage, S. Deep neural networks effective at learning high- dimensional Hilbert-valued functions from limited data. In Bruna, J., Hesthaven, J., and Zdeborova, L. (eds.), Pro- ceedings of the 2nd Mathematical and Scientific Machine Learning Conference, volume 145 Proceedings of Ma- chine Learning Research, pp. 1-36. PMLR, 16-19 19 Aug 2022. https: / http///ooceedieeeocce ngs mlr ..l.press/////// v1 4. htt.5//dccc222a..2. html.
Davi, C. and Braga-Neto, U. PSO-PINN: Physics-informed neural networks trained with particle swarm optimization. arXiv:2202.0199 [physics], February 2022.
Daw, A., Bu, J., Wang. S., Perdikaris, P., and Karpatne, Mitigating propagation failures in physics-informed neu- ral networks using resiin-reeamlle-eeeeee (13) sampling. In Proceedings of the 40th International Conference on Machine Learning, pp. 7264-7302, 2023.
press/
Avron, H., Maymounkov, P., and Toledo, S. Blendenpik: Supercharging LAPACK's least-squares solver. SIAM Journal Scientific Computing, 32(3):1217-1236, 2010.
Fletcher, R. Practical Methods of Optimization. John Wiley & Sons, New York, 2nd edition, 1987. ISBN 978-0-471- 91547-8.
Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky,, P., Budden, D., Cai, T., Clark, A., Dani- helka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C.,, Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kap- turowski, S., Keck, T., Kemaev, I., King, M., Kunesch, M.. Martens, L., Merzic, H., Mikulik, V., Norman, T., Papamakarios, G., Quan, J., Ring, R., Ruiz, F., Sanchez, A., Schneider, R., Sezener, E., Spencer, S., Srinivasan, S., Stokowiec. W., Wang, L., Zhou, G., and Viola, F. The DeepMind JAX Ecosystem, 2020.
Fridovich-Kens Keil, S., Gontijo Lopes, R., and Roelofs, R. Spec- tral bias in practice: The role of function frequency in generalization. Advances in Neural Information Process- ing Systems, 35:7368-7382, 2022.
Gower, R., Kovalev, D., Lieder, F,, and Richtárik, RSN: randomized subspace Newton. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Halko, N., Martinsson, P. G., and Tropp. A. Finding
Baker, N., Alexander, F., Bremer, T., Hagberg, A., Kevrekidis, Y., Najm, H., Parashar, M., Patra, A., Sethian, J., Wild, S. M., and Willcox, K. Workshop report on ba- sic research needs for scientific machine learning: Core technologies for artificial intelligence. Technical re- port, USDOE Office of Science (SC), Washington, DC (United States), 2019. URL https: wwt.wwww.......o////////. gov/servles pptpppuu11111/
structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions, SIAM Review, 53(2):117-288, January 2011. ISSN 0036-1445, 1095-7200. doi: 10. 7,.11370907778066
Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P., Wang, S., and Yang. L. Physics- -nyiics-informed machine learning. Nature Reviews Physics, 3(6):422-440, 2021. doi: 10. 020004422544002-000333
Belkin, M., Hsu, D., Ma, S., and Mandal, Reconciling modern machine-lcarning practice and the classical bias- variance trade-off. Proceedings of the National Academy of Sciencess (5849)2:1544-15888 2019.
Kidger, P. and Garcia, C. Equinox: Neural networks in JAX via callable PyTrees and filtered transformations, October 2021.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. International Conference on Learning Rep- resentations (ICLR), 2015. pttps://arxivve abs/1412. 6980.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. Leary, C., Maclauin,, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Mille S., and Zhang, Q. JAX: Composable transformations of Python+NumPy programs, 2018.
König, H. and Richter, S. Eigenvalues of integral operators defined by analytic kernels. Mathematische Nachrichten, 199(1):441-155 1984. ISSN 0025584X, 15222616. doi: 10. 10.1002//ana. 1tt00//maaa.1941199909 19841190113.
Chang, X., Li, Y., Oymak, S., and Thrampoulidis, C. Prov- able benefits of overparameterization in model compres- sion: From double descent to pruning neural networks. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 35, pp. 6974 6983, 2021.
Krishnapriyan, A., Gholami, A., Zhe, S., Kirby, R., and Mahoney, M. W. Characterizing possible failure modes in physicc--ifooee informed neural networks. Advances in neural information processing systems, 34:26548 225660, 2021.
Ciarlet, P. G. The Finite Element Method for Elliptic Prob- lems. North-Holland, 1978.
10

---

Near-optimal Sketchy Natural Gradients for PINNS
Li, H., Xu, Z., Taylor, G., Studer, C., and Goldstein, T. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
van der Meer, R., Oosterlee, C. W., and Borovykh, A. Op- timally weighted loss functions for solving PDEs with Neural Networks. Journal of Computational and Applied Mathematics, 405:1 13887, May 2022. ISSN 0377-0427 doi: 10.1016/j.cam camp///ccam.200000 2021.022.1111 13887.
Müller, J. and Zeinhofer, M. Achieving high accuracy with PINNs via energy natural gradient descent. In Inter- national Conference on Machine Learning, pp. 25471- 25485, PMLR, 2023.
Wang. S., Teng, Y., and Perdikaris, P. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal on Scientific Computing, 43(5):A3055AA33 January 2021. ISSN 1064-8275. doi: 10.1 1137/20M1388/44
Nurbekyan, L., Lei, W., and Yang, Y. Efficient natural gradi- ent descent methods for large-scale PDE-based optimiza- tion problems. SIAM Journal on Scientific Computing, 45(4):A1611-A165, 2023.
Wang, S., Yu, X., and Perdikaris. P. When and why PINNs fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449: 10768, January 2022. ISSN 0021-9991. doi: jep.01//..pp.222..1100
Rahaman, N., Baratin, A., Arpit, D., Draxler, F., Lin, M., Hamprecht, F., Bengio, Y., and Courville, A. On the spec- tral bias of neural networks. In International conference on machine learning, pp. 5301-5310. PMLR, 2019.
Wang, S., Sankaran, S., and Perdikaris, P. Respect- ing causality for training physics- informed neural net- works. Computer Methods in Applied Mechanics and Engineering, 421:116813,, 2024. doi: 10.1016/j.cm. 2024. 1 11633. URL https: /tg///.o.ooo...oo/////// 1001// i. cma. 2024 116813.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, February 2019, ISSN 0021-9991. doi: /0p.///.ccp..01.
Weinan, E. et al. The dawning of new era in applied mathematics. Notices of the American Mathematical Society, 68(4):565 2021.
Rathore, P.. Lei, W., Frangella, Z., Lu, L., and Udell, M. Challenges in training PINNs: A loss landscape perspec- tive. In Salakhutdinov, R., Kolter, Z.,, Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the s1 International Conference on Ma- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 42159 42191. PMLR, 21-27 Jul 2024. StRtsttp:: /ttp://ppoccceees/e/eeeee m.r.prress/ 0235/rathore220 html..
Wong, J. C., Ooi, C. C., Gupta, A., and Ong. Y.-S. Learning in sinusoidal spaces with physics-informed neural net- works. IEEE Transactions on Artificial Intelligence, 5(3): 985-1000, 2022.
Wu, C., Zhu, M., Tan, Q., Kartha, Y., and Lu, L. A comprehensive study of non-adaptive and -siduull-based adaptive sampling for physics-informet neural networks. Computer Methods in Applied Mechanics and Engineer- ing. 403: 1 15671, 1, January 2023. ISSN 0045-7825. doi: 10.10166j.cma. 202.. 115671.
Rohrhofer, F. M., Posch, S., Gößnitzer, C., and Geiger, B. C. On the role of fixed points of dynamical systems in training physics-informed neural networks. Transactions Machine Learning Research, 2022. URL https: //openreview net/ forum d= 1dd55ccTTzgg5.....
Yang, M., Xu, D., Wen, Z., Chen, M., and Xu, P. Sketch- based empirical natural gradient methods for deep learn- ing. Journal of Scientific Computing, 92(3):94, 2022.
Saibaba, A. K., Lee, J., and Kitanidis, P. K. Randomized al- gorithms for generalized Hermitian eigenvalue problems with application to computing Karhunen-LLovve expan- sion. Numerical Linear Algebra with Applications, 23(2): 314-339, 2016. ISSN 1099-1506, doi: 10. 1026/
Zeng, Q. Bryngelson, S. H., and Schaefer, F. T. Competitive physics informed networks. In ICLR 2022 Workshop on Gamification and Multiagent Solutions, April 2022.
Sejnowski, T. J. The unreasonable effectiveness of deep learning in artificial intelligence. Proceedings of the Na- tional Academy of Sciences, 117(48):30033- 30038, 2020.
Zhang, S., Zhao, H. Zhong, Y., and Zhou, H. Why shal- low networks struggle with approximating and learning high frequency: A numerical study. (arXiv:2306. 13311.. November 2023. doi: 10.488550/axxiv.23366 17301
Trefethen, L. N. and Bau, D. Numerical Linear Algebra. SIAM, 1997. ISBN 0898713617.
Zhu, C.,, Byrd, R. H., Lu, P., and Nocedal, J. Algorithm 778: L-BFGS-B: Fortran subroutines for large-scale bound- constrained optimization. ACM Trans. Math. Sofnw, 23 (4):550-560, December 1997. ISSN 0098-3500. doi: 10. 10.1145//77923.277933.
Tropp, J. A., Yurtsever, A., Udell, M., and Cevher, V. Practi- cal sketching algorithms for low-rank matrix approxima- tion. SIAM Journal on Matrix Analysis and Applications, 2017.

---

Near-optimal Sketchy Natural Gradients for PINNS
A. Natural gradients for two simple problems consider two basic problems
[ (u(x) - f(x))) min u
(14)
and
[ [ dx) (x) - dr
min 17
(15)
where I,2 is the space of measurable, square integrable functions and H. is the space of functions u whose derivatives are square integrable and u(-1) = 0. For details on these spaces and some of the results used below, see the classic text (Adams, 1975). both cases, the data function given. For completeness, we show some results that are standard the finite element literature (Ciarlet, 1978).
Problem (14) is basic data fitting problem. Here, have assumed infinite training data and access every function in the space L In practice, limited data is available for training, the space of allowable approximating functions is finite dimensional, and the integral is replaced by quadrature rule on the data. However, in the following theoretical discussion we will keep the infinite data assumption. Problem (15) solves an ODE boundary value problem. The minimizer satisfies
d²u
dx?
=
with boundary conditions u(-1) = and u' u1(1) 0.
consider now the minimization problems with functions u restricted to linear, dimensional subspace S, with basis functions (z)). The optimal solution U, S to problem (14) satisfies the linear system
MU.
where U is the vector of coefficients of the optimal u(r):
u(x) Σ Unhi(z),
M is the symmetric, positive definite matrix with entries
Ma.1 [- (s(00)0
and F is the vector
F2 b [" f())
These are standard results from the finite element literature (Ciarlet, 1978) where M called the mass matrix. We consider
OUU) = [ Σ - ffx))))))((())))
N
2
dr
and consider finding the optimal solution U, via gradient descent (GD). The gradient of is
MU-F
By considering the error E U U. we can write gradient descent terms of the error with direction
ME
and with this standard trick we can neglect the RHS f and F and consider U.
12

---

Near-optimal Sketchy Natural Gradients for PINNS
Figure 5. Contours of well-conditioned energy landscape (left) such achieved for data fitting problem with finite element basis. poorly behaved energy landscape (right) in which the steep vertical gradients limit the learning rate and lead to slow convergence in the horizontal modes. For simple Neural Networks applied a data fitting problem, is demonstrated that the slowly converging modes are oscillatory (high wavenumber). The use of the Gram Matrix in the NEG method in an idcalized setting moves the landscape from the right the left. our practical implementation. it promotes convergence finite number of modes.
We consider now whether M is well conditioned (Trefethen & Bau, 1997) and the impact this has on the convergence of GD. We calculate
s(U) - [[] ( N (
MU,
where the LHS (:: -) the Fuclidean inner product and see that the matrix M relates the 1,2 metric (energy) of the underlying problem to the Euclidean metric of the subspace basis coefficients. We can consider the orthonormal basis of eigenvectors of M, {V3} with corresponding eigenvalues A 0. If we write
N
U
V
j=1
then the gradient has the simple form
N MU
Σ
F=1
and the level sets of energy MU, simplify to ellipsoids aligned with the coordinate axes
N E = Σ 1,,))
The gradient perpendicular the level sets, which gives graphical addition the ideas below.
If } is orthonormal set in 2 then M is the identity matrix, and GD iterations will tend straight the optimal solution U. = 0, as shown in Figure 5 (left). The condition number of M is defined as
<(M) Amax/ Amin-
If M has small condition number, GD convergence is almost uniform in energy. M has a large condition number, then the learning rate for GD limited by the large eigenvalues and the small eigenvalue components converge slowly, see Figure (right). It is well known that PINNs and neural networks, general, suffer from spectral bias; training proceeds slowly in wavenumber (oscillatory) components. We show this explicitly in Section B below, following the ideas in (Zhang et al., 2023).
The linear subspace discussion has relevance to nonlinear approximations u(@) such as approximations by neural networks if consider the tangent spaces of these approximations.
Remark A.1.
13

---

Near-optimal Sketchy Natural Gradients for PINNS
We now describe the purpose of the natural energy update in Equation (5). For the linear subspace data fitting problem above, parameters Og are just the basis coefficients Ug, and
VoL = MU
where we have again neglected the term from the data function f. The operator F in Equation (2) is the identity operator, in this case, and 0ug/00, is the basis function (z)). Thus, the Gram matrix G is the mass matrix M, and the natural energy gradient is
MU.
If the exact M were used, would obtain the identity operator, with uniform convergence in the span of S. By using an efficiently computed, approximate M1, we are preconditioning GD to increase the convergence rate of some components of the solution. In the case of neural networks and PINNs, we are increasing the convergence rate of medium wave-number components of the solution.
For completeness and use the next section, we consider the ODE problem (14). Here, we look for an optimal solution in the same subspace, and again neglecting the data f, we have gradient KU where
Ki. = ["= days duy dx, dæ da
known as the stiffness matrix in the FE community. We reuse the previous notation and write
((((e)) dz
C
N
di/s
2
((U) =
U
= -U,
-U,KU)
de
and see that the matrix K relates the natural (Dirichlet) energy
[ du
dr
to the Euclidean coefficient metric. If the condition number of K is large, then convergence in the small eigenvalue components will be slow in GD iterations. Following the same approach above, the natural energy gradient for this problem is
KtKU
and efficiently calculated approximate K"t improves convergence in of the components.
B. Condition number of the Gramian for Neural Networks and PINNs
We consider a specific subspace S and two different bases for it. We take to be the set of piecewise linear functions U(x) on N equal subintervals of length = 2/N of x E [-1, 1] with U(0) = We can represent (xx) with a standard FEM basis (x)) (i = I.... N), which has value one at one grid point v; -1 ih and zero at other grid points. We could also represent (z) with a basis (x) (i O....N 1) given by
(x(x) = o xz))
where a is the ReLU activation function. It is shown in (Zhang et al., 2023) that there is nice reduction from a general lne-layer neural network that leads to this expression. In practice, the positions 2; vary in the optimization; however, following (Zhang et al., 2023), for the sake of tractability, keep them fixed and continue in the linear subspace framework. We can write (x) e S in either basis with different coefficients
Ulz)= Uu(r(()= Σ Vioi(z).
The optimal solution can similarly be expressed in either basis as
MU F and M2V G.
(16)
14

---

Near-optimal Sketchy Natural Gradients for PINNS
where
S dibydo
M1.9
L Φ d
M2.33
Fj J f
(17)
[ (z)
Gj
(18)
Direct calculation shows that
D203 = lug and D2U; = IV,
interior grid points where
D20; == 0- 2Φj + Φj+1.
There modifications at the boundary grid points that will not be careful with here. With this result, we from (17 and 18) that
F=hD2
and from (16)
-
from which can be read
M1 = D2M2D2
or M2 = D2 ' is known that Mj is well conditioned. An explicit computation using discrete Fourier transform techniques gives
<(M) 3
for all N. The operator D2 is similarly known to have condition number N2 with large eigenvalues associated with large wavenumber components. Thus, M2 has condition number O(N4) and has slow convergence for these oscillatory components in an GD optimizer. The O(N+) condition number for this reduced NN is shown more generally in (Zhang et al., 2023).
We can apply the same approach to the ODE boundary value problem using stiffness matrices K1 and Kg leading with the same algebra to
K₂ 1==Dz' K, D2
1
Since K; = - D22/h, we have
N K₂
DJ
2
somewhat unexpected result. Careful handling of the boundary terms shows that the exact result above applies with modified
D200 = C1
and
D2CN-1 = 2CN-1 CN-2
The result above gives evidence that the Gram matrix for the PINN approximation of the model problem (15) does have large condition number, and that the use of the approximate natural energy gradient will promote the training of oscillatory modes the solution. The work in (König & Richter, 1984) gives some evidence that the exponential decay in the cigenvalues of a PINN Gram matrix observed computationally is expected when an analytic activation function such as swish used.
15


Reviews:
[{"review_id": "90Kq2FOWOT", "reviewer": "Reviewer_pKbx", "summary": "This paper introduces a novel natural gradient descent method based on sketching. Instead of computing the search direction of the natural gradient exactly, the paper uses the sketching technique to compress the large Gram matrix into a smaller one with a random Gaussian matrix. The algorithm is interesting and theoretically sound. The theoretical analysis shows that the accuracy of the sketching approximation depends on the least eigenvalues of the Gram matrices, which typically exhibit eigenvalue decay. Additionally, the paper also conducts comprehensive experiments on comparing the proposed methods with ADAM and BFGS.  Overall, this is a very interesting study!\n\n## update after rebuttal\nI maintain my positive score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to carefully review our manuscript and appreciate the positive review of our work.\n\n**Questions For Authors:**\n\nThe proposed algorithm is, indeed, not limited to training PINNs. We began the paper by focusing specifically on improving the training of PINNs, and thus, all of our experiments are examples of PINNs. Our algorithm’s benefits depend on the matrix $\\mathbf{G}$ structure, which exhibits exponential decay. This property is related to the spectral bias of neural networks and not explicitly tied to PINNs or the underlying PDE. However, the structure of PINNs tends to worsen spectral bias in ways that can be related to the underlying PDE [1].  In seeking to understand the performance of our algorithm, we discovered that our method may be more broadly applicable. We hope to explore this in future work. \n\nAs a simple illustrative example, consider training a standard feed-forward fully connected neural network on a set of features and labels using mean squared error as the loss function. For this example, a different but related Gramian matrix $\\mathbf{G}$ can be used to perform a “pre-conditioned” gradient descent (see appendix A for a construction of this matrix in a simple setting.) This matrix should have exponentially decaying eigenvalues for analytic activation functions, and one could use Algorithm 1 for this task. \n\n[1] Bonfanti, Andrea, Giuseppe Bruno, and Cristina Cipriani. \"The challenges of the nonlinear regime for physics-informed neural networks.\" Advances in Neural Information Processing Systems 37 (2024): 41852-41881.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes. The natural gradient descent is not commonly adopted for training PINNs, mainly because of its high computational cost for the search direction (Gram matrices). This paper alternatively computes an approximate natural gradient direction by sketching. This technique, although is not new, should be effective and meaningful.", "theoretical_claims": "Yes. The proof supports the effectiveness of the sketching technique, under the condition that the eigenvalues of Gram matrices are decaying.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The algorithm proposed is not limited to training PINNs. I think it can be applied to broader deep learning applications.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "The algorithm proposed is not limited to training PINNs. I think it can be applied to broader deep learning applications. Do I understand correctly? Because the design of the neural network does not depend on the structures of PDEs or PINNs. So based on this, why does the paper only focus on applying it to PINNs? And also the title of the paper. Why is the method restricted to PINNs?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OLkQFSE3kM", "reviewer": "Reviewer_tL44", "summary": "The manuscript discusses the application of randomized numerical linear algebra to scale natural gradient methods for the training of physics informed neural networks (PINNs). More precisely, a randomized eigensolver is employed to solve the linear system in the natural gradient algorithm at every step. The authors demonstrate that this can be done efficiently and scalable, i.e., without the need to assemble the Gramian matrix. It is expected that this realization helps the widespread adoption of natural gradient methods for the training of PINNs.\n\n## update after rebuttal\nI maintain my positive evaluation.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our manuscript and for their helpful suggestions, comments and questions. We greatly appreciate their positive assessment of our work. \n\n**Relation To Broader Scientific Literature:**\n\nWe thank the reviewer for the suggestion and have included a brief discussion of this reference in the manuscript in the “related work” section.\n\n**Weaknesses:**\n\nThank you for the suggestion. We will include some numerical experiments in a new appendix. As a preliminary answer: the low-rank structure of the Gramian relates to the phenomenon of Spectral Bias. Small eigenvalues of the Gramian correspond to high-frequency components that neural networks struggle to learn. The literature on spectral bias suggests that higher-frequency components need more sample points for effective learning than lower-frequency ones and that the spectrum of the Gram matrix may depend on the density of samples near these higher-frequency components. \n\n**Questions For Authors:**\n\n1. Damping can be included in the solution approach, and we thank the reviewer for the suggestion. We will add it to future work. \n2. The sketch size at each iteration of training depends on a fixed “over-sketch” parameter $p$ and on the computed rank of the sketched matrix $GM$ at the previous iteration of training. The sketch size is $p+r$ where $r$ is updated at each iteration of training. This means that the sketch size is adaptively growing to incorporate the growing rank of $G$ as training progresses. We thank the reviewer for this question and have added some explanation in the manuscript that makes this clearer and explains how the initial rank is estimated to initialize Algorithm 1. \n3. \n   -  We agree that Nyström methods are a natural choice for further exploration. We appreciate the comment and hope to explore the benefits of Nyström methods in future work. As an additional comment, for Nyström methods that rely on repeated computation of a matrix-vector product, the computational cost of the nested auto-differentiation in $G$ is a major impediment. See the fourth comment under “Claims and evidence” in the response to reviewer nY1h for a more detailed discussion of this last point.\n    -  We elected to use one/two pass results from the classical randomized linear algebra literature for a couple of reasons. (1) First, to the best of our knowledge, error bounds for Nystrom methods are more subtle and depend on how columns are sampled. For our use case where we have a PSD matrix with exponentially decaying eigenvalues, error bounds for one/two-pass eigenvalue approximation show that by controlling the tolerance, we can drive the error close to machine precision. Second, because we can efficiently compute the sketch $ GM $ via equation (7), our manuscript already demonstrates the value of sketching for natural gradients in terms of both accuracy and scalability, with the added advantage of using a method that is easy to understand and implement. \n4. We thank the reviewer for the helpful suggestion and have updated the manuscript to provide plots of the loss/errors during training to showcase the convergence.", "claims_and_evidence": "The claims made in the manuscript are properly substantiated. The numerical experiments are convincing.", "methods_and_evaluation_criteria": "The setup of the numerical experiments is reasonable.", "theoretical_claims": "I did not check the theoretical result concerning the low rank of a two layer network.", "experimental_designs_or_analyses": "I checked the setup of all numerical experiments and they are reasonable.", "relation_to_broader_scientific_literature": "The relevant literature is appropriately discussed. To the best of my knowledge, no relevant work is omitted from the discussion in the literature. The authors may however consider to discuss https://arxiv.org/pdf/2402.01868, which uses a Nyström method to design a preconditioner for a CG in Newton's method. As ENGD is different than Newton (and typically much more effective) I see no direct implications for the authors work.", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "Weaknesses:\n- How does the low-rank structure of the Gramian change with varying sample size? Can the authors provide numerical experiments for this?\n\nStrengths:\n- The numerical results are convincing. It is very nice that a network with a million parameters can be trained in only a handful of minutes.", "other_comments_or_suggestions": "See below.", "questions_for_authors": "1. It is common practice to include damping into natural gradient methods, i.e., adding a scaled identity on top of the Gramian matrix. Can this be included in the solution approach?\n2. The rank of the Gramian seems to grow with training iterations. Does this require larger sketch sizes in later stages of training?\n3. Can the authors comment on the choice of randomized method? Why did you not employ Nyströms method which is taylored for psd matrices?\n4. Can the authors provide some loss/error plots including statistics to showcase convergence of the optimizer?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H45lusaNeN", "reviewer": "Reviewer_D1jG", "summary": "This work improves the computational efficiency and estimation accuracy of NEGD by leveraging the structural properties of the Gram matrix and introducing the classical RSVD method. It effectively addresses the high storage and computational costs associated with the Gram matrix while also enhancing the neural network’s ability to learn high-frequency components.\n\n## Update after rebuttal \n\nThe authors claim their contribution lies in analyzing the matrix structure and linking it to known error bounds from randomized numerical linear algebra, stating this is novel in the context of natural gradients for PINNs.   However, from my perspective, the proposed method is essentially a straightforward application of RSVD  to NEGD, and the spectral error bounds discussed are drawn from existing literature. Based on both the paper and the authors’ rebuttal,  I do not see a clear advantage or new theoretical insight specific to the combination of RSVD and PINNs. Therefore, I find the theoretical contribution to PINNs limited and maintain my original scores.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful and detailed suggestions and comments and the time and effort they put into reviewing our manuscript. \n\n**Experimental Designs Or Analyses**\n\nWe will add comparisons with ADAM to 5.2-5.4. \n\nAt the end of section 5, we specify how $p$ is chosen and set the tolerance to 1E-13 across all experiments. \n\nThe sketching size changes at each iteration of training adaptively as the rank grows during training as a result of the underlying parameter space, and therefore, we cannot directly compare sketching sizes for different architectures.\n\n**Essential References Not Discussed**\n\nRegarding “Streaming …” please see the comment (2) under Weaknesses below. \n\nRegarding RSN: We agree with the reviewer that it would be interesting to compare RSN to SNGD, but we believe the computational complexity of SNGD to be either similar or lower than RSN. Both RSN and SNGD require a randomized approximation: RSN requires the Hessian of the loss at every iteration of training, and SNGD requires the matrix $G$ at every iteration of training.  The Hessian entails computing the second derivatives w.r.t. the parameters of the loss function, which depends on the PDE residual, and $G$ requires computing the first derivatives w.r.t. the parameters of the PDE residual itself. Moreover, for SNGD, the sketch $GM$ can be computed efficiently via equation (7). \n\n**Weaknesses**\n1) The main innovation is that we identify and exploit spectral structure inherent to the matrix $G$ to develop a fast, memory-efficient, and accurate natural gradient method. The structure of $G$ is the best-case scenario for the approximation capacity of RSVD, due to the exponential decay of the eigenvalues, which guarantees error near machine precision. Previously, NEGD was limited to only very small neural networks. Therefore, our work represents a significant advance. We will consider more sophisticated estimation techniques in future work. \n\n2) We don't believe our use case falls into the streaming framework: gradient descent iterations do not follow a known linear update (see more details in the next paragraph). We employ the methods recommended by the same authors in the related paper “Fixed-Rank Approximation of a Positive-Semidefinite Matrix from Streaming Data”, designed specifically for PSD matrices (such as $\\mathbf{G}$). On page 3, they state that in the absence of constraints such as streaming, they recommend the general-purpose methods we use.\n\nConcretely, to clarify why updates to $G$ do not follow a known linear map, notice that $\\mathbf{G}\\_{i+1}$ depends on $\\left(  \\partial{\\theta\\_k} \\circ \\mathcal{D}\\right) \n\\left[ u\\_{\\theta}(\\theta\\_{i+1})\\right]$ through equation (6). We can write $u\\_\\theta(\\theta\\_{i+1}) = u\\_\\theta\\left( \\theta\\_i - \\eta \\mathbf{G}^{\\dagger}(\\theta\\_i)\\nabla\\_{\\theta}L(\\theta\\_i) \\right)$. Because $u\\_\\theta$ is non-linear, the relationship between $\\mathbf{G}\\_{i+1}$ and $\\mathbf{G}\\_i$ is non-linear. \nWe are not aware of a way to express $\\mathbf{G}\\_{i+1}$ as $\\mathbf{G}\\_{i+1} =  \\eta\\mathbf{G}_i+ \\nu \\mathbf{H}$. \n\n**Other Comments Or Suggestions**\nWe thank the reviewer and have updated the manuscript to fix each error. \n\n**Questions For Authors**\n1. We have fixed a typo in Algorithm 1: $Q \\gets Q[:,1:(p+r)]$,  so $T = Q^{T} G Q$. \n2. Lemma 3.1 shows that for a 1D, 1-layer network with analytic activation functions, the eigenvalues decay exponentially. We agree that the discussion of how to choose the threshold tol should be clarified and tied more directly to the error bounds and Lemma 3.1. Based on the exponential decay supported by Lemma 3.1 and by the empirical evidence,  the tolerance is set s.t. $\\lambda\\_{j+1}$ is close to $10^{-16}$. \n3. We agree that this point should be clearer. Because the error in equation 8 is bounded by $C(p,r,e) \\lambda\\_{r+1}$, it depends on the oversampling parameter only through a polynomial constant. When $r$ is sufficiently large, the exponential decay of the eigenvalues drives the error close to machine epsilon, and the oversampling parameter p does not make much of a difference. We have clarified this point in the main text.\n4.  We have added some discussion in the last paragraph of section 4.\n        1. Because the matrix $\\mathbf{G}$ is ill-conditioned, computing the pseudo inverse via least squares introduces computational instability that is ameliorated by cutting off small near-zero eigenvalues. \n        2. In section 5.4, we compare the rank above the cut-off of 1e-13 of two-pass SNGD and the original ENGD and see that SNGD is able to find more directions in parameter space above this cut-off. This suggests that SNGD is finding a better and flatter local minimum,  which is a phenomenon that has been observed to be useful for generalization and related to spectral bias. Finally, the randomness introduced by multiplying by a Gaussian matrix at each training iteration and the mixing effect of Guassians contribute to the enhanced performance.", "claims_and_evidence": "The core argument of this paper is that introducing randomized sketching techniques for PINNs optimization can enhance both the computational efficiency and accuracy of the NEGD method. This claim is validated through extensive numerical experiments, where the proposed method is compared against ADAM, BFGS, and ENGD.\n\nHowever, one concern remains: Sections 5.2–5.4 present experimental results showing that SNGD outperforms BFGS and ENGD in both efficiency and accuracy. Why is there no comparison with the ADAM method in these sections?", "methods_and_evaluation_criteria": "Yes, this work proposes a new method that addresses the computational bottleneck in PINNs and improves estimation accuracy.", "theoretical_claims": "This paper does not provide  novel theoretical results and does not verify any proofs.", "experimental_designs_or_analyses": "Yes, I have reviewed the experimental content and results in Section 5, and overall, they appear reasonable and effective. However, the following concerns remain: \n1. Section 5 demonstrates SNGD's superior performance over NEGD, ADAM, and BFGS across multiple examples. While the paper acknowledges the impact of $tol$, p on algorithm performance, it does not specify their values for each experiment.\nProviding these parameter values would enhance the reproducibility of the results and strengthen the credibility of the findings.\n\n2. Table 1 shows that both SNGD and ADAM achieve similar accuracy on the 5,000-parameter (A2) and 1,000,000-parameter (A3) networks.  However, this does not fully demonstrate SNGD's scalability, as it is unclear whether A2 and A3 use the same p+r. If they differ (e.g., if A3 uses a larger p+r, meaning the sketching matrix M has a higher dimension), this conclusion may lack sufficient credibility. Could you provide a more detailed explanation?", "relation_to_broader_scientific_literature": "The paper introduces the advantages and challenges of PINNs, discusses the research progress of commonly used solution algorithms, and analyzes their strengths and weaknesses. Among them, NEGD (Müller & Zeinhofer, 2023) currently achieves the highest estimation accuracy. However, NEGD suffers from high computational complexity.\n\nTo address this, the authors employ single-pass and two-pass sketching techniques from Saibaba et al. (2016) to improve both the computational efficiency and estimation accuracy of NEGD, proposing a new method called SNGD.", "essential_references_not_discussed": "The paper \"Streaming low-rank matrix approximation with an application to scientific simulation\" introduces a more efficient single-pass RSVD algorithm than Saibaba et al. (2016), yet this work does not mention it. Given that the study employs RSVD to accelerate NEGD computation, why not adopt the more efficient single-pass RSVD from that work? Could you clarify this choice?\n\n\n\nAnother common approach for accelerating over-parameterized iterative estimation using randomized sketching techniques is RSN (see RSN: Randomized Subspace Newton). Given the iterative estimation structure used in this paper, RSN may have a lower computational complexity than RSVD. In addition to utilizing dense Gaussian matrices, RSN allows the use of sparse or orthogonal sketches, which could offer additional computational benefits.\nTo further strengthen the contributions of the paper, I suggest comparing the RSN method with the RSVD-based approach introduced in this work. A thorough comparison could better illustrate the advantages of using RSVD for gradient descent acceleration and preconditioning, particularly from the perspective of spectral bias analysis. This would enhance the practical significance and applicability of the conclusions presented in the paper.", "other_strengths_and_weaknesses": "Strengths\n\nThe paper proposes the SNGD method by incorporating randomized sketching techniques, which improve both the computational efficiency and estimation accuracy of NEGD.\n\nWeaknesses\n\n1） However, the SNGD approach primarily relies on a direct application of the classical randomized SVD (RSVD) algorithm for low-rank approximation of the Gram matrix. While this method is straightforward, the paper does not fully demonstrate the advantages of integrating RSVD with NEGD compared to other estimation techniques. The methodological innovation appears somewhat limited.  Could you provide deeper insights and conclusions about the SNGD method, beyond those merely derived from existing low-rank approximation results?\n\n2） This work does not  adopt the more efficient single-pass RSVD from  \"Streaming low-rank matrix approximation with an application to scientific simulation\"  to accelerate NEGD computation.", "other_comments_or_suggestions": "1.\tline 68:  The first occurrence of NEGD should follow the standard convention of writing the full term followed by the abbreviation in parentheses.\n\n2. line 226: $G^{-1}_{(p+r)}$-> $ \\hat G^{\\dagger}_{(p+r)}$\n\n3.\tlines 298-299: Ensure consistency in the naming of methods in experimental comparisons:  In Table 1, update \"ONE\"-> \"SINGLE-PASS\" and \"TWO\" -> \"TWO-PASS\".\n\n4.\tline 175:  There is an error in Algorithm 1, specifically in the single-pass step. The formula $ (QM)^{-1}$ is incorrect as the dimensions do not match.\n\n5.\tPlease ensure uniform formatting of references, particularly in paper titles, where only the first word should be capitalized, and the rest should be in lowercase.", "questions_for_authors": "1. Does Equation (7) have an extended version specifically for the two-pass method mentioned in Algorithm 1? Equation (7) cannot be directly used for the two-pass update:  $T <- Q^GQ $\n\n2.The approximation $(\\sum^n_{j=r+1}\\lambda^2_j)^{1/2}\\approx \\lambda_{ r+1}$ and the error bound  $E\\|G-\\hatG_{p+r}\\|\\lesssim10^{-16}$ hold only under the assumption that tol is set to 10^{-16}. However, the paper does not provide a theoretical guarantee that tol should be precisely  10^{-16}. Currently, this assumption is only supported by the experimental results in Figure 1, which may not be sufficiently convincing. Could you provide a more detailed explanation of how the threshold tol is determined? A more rigorous justification would strengthen the validity of the two formulas above.\n\n3. Section 5 claims that Section 4.1 explains why the training process is much less sensitive to the oversampling parameter p, but this explanation is unclear. Could you clarify this point more directly? \n\n4. This paper applies randomized SVD to approximate the Gram matrix, improving computational efficiency. Section 5 shows that SNGD, using \\hat G, outperforms NEGD with exact G in both efficiency and accuracy. Could you clarify in the main text why randomized SVD enhances NEGD’s estimation accuracy? A stronger explanation would reinforce the experimental conclusions and highlight the paper’s contributions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DfgrFzmgfj", "reviewer": "Reviewer_nY1h", "summary": "This paper proposes SNGD, a sketched version of natural gradient descent (NGD), for training PINNs. SNGD uses sketching to scale a previous optimizer, ENGD, to larger neural networks for PINN training. The authors evaluate SNGD against Adam, BFGS, and NEGD on several benchmark problems and demonstrate that SNGD outperforms the competition.\n\n## update after rebuttal\nI thank the authors for responding to my questions and comments on their paper. However, I still believe that my concerns about the experiments have not been adequately addressed. For example, the authors say that they can compare to the results in “Challenges in Training PINNs: A Loss Landscape Perspective”, but this might not be a good comparison due to differences in network architecture, hyperparameters, etc. I am inclined to maintain my score at this time.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the many helpful suggestions, for their careful and detailed comments and for the positive comment about the novelty and utility of the work. \n\n**Claims And Evidence:**\n\n- If we understand your comment correctly, we believe that our observation about rank deficiency is completely in line with your observation about discretization. Indeed, the expression for (6) describes the continuous form of the matrix before discretization, while $\\mathbf{G}$ refers to the matrix that we get under an appropriate discretization of (6). We use the convention of denoting continuous objects with calligraphic font and their discretized counterparts as bolded objects. We have added a sentence to clarify this notational convention. \n \n- At the top of section 3, we state that Figure 1 corresponds to an experiment from section 5.4. This info has been added to the caption. The exponential decay of the eigenvalues of $\\mathbf{G}$ is a feature that we observed across all of the experiments presented in the paper. This is briefly mentioned in section 3. Given your comment, we added a sentence stating this at the start of section 3 to clarify this. \n\n- The constant $C(p,e,r)$ is polynomial in $r$. Our case has the ideal feature of exponential decay in the eigenvalues, which means that for  $r$ sufficiently large, the decay is rapid enough to beat the polynomial growth of $C$. We have added a sentence at the top of page 5 to explain this. \n\n- In our setting, $\\mathbf{G}$ is dense with $|\\theta|^2$ entries, and computing its matrix-vector products is prohibitively expensive. These require computing nested auto-diffs of the residual w.r.t. each parameter. This would need to be computed anew in every CG iteration. Therefore, we opt for directly sketching $\\mathbf{G}$. In “Randomized Nystrom Preconditioning” the authors explicitly caution that in order for Nystrom preconditioning to be effective, matvecs must be reasonable in terms of computational cost. However, we speculate that a low-rank preconditioner might have, in some ways, a similar effect to the benefits gained by cutting off the smallest eigenvalues via sketching. \n\n- We really appreciate this question and thank the reviewer. We will update the manuscript to include this information. The memory requirement for NEGD is $\\mathcal{O}(|\\theta|^3)$ and for SNGD it is $\\mathcal{O}(|\\theta|(p+r))$. Focusing on the computational cost of the solve per iteration, least squares costs  $\\mathcal{O}(|\\theta|^3)$  for ENGD, whereas Algorithm 1 has a cost $\\mathcal{O}(|\\theta|(p+r)^2)$ when $|\\theta|$ is large enough.   \n\n- The error bound is “near-optimal” in the sense that due to the exponential decay of the eigenvalues of $\\mathbf{G}$, we can pick a tolerance that guarantees that the error is close to machine precision. We feel that given the nature of the decay, it is fair to characterize the error in this way. \n\n**Methods And Evaluation Criteria:**\n\nWe thank the reviewer for this comment and agree with their assessment. We are working on a benchmark for a more challenging problem and appreciate the reviewer’s suggested additional benchmark problem. We will do our best to update the manuscript to include both of these.\n\n**Theoretical Claims:**\n\nThe decay rate for ReLU is polynomial; however, when using analytic activation functions the decay rate becomes exponential. Please see the discussion near Lemma 3.1. \n\n**Experimental Designs Or Analyses:**\n\n- Because ADAM is a first-order optimizer, it cannot achieve the high level of accuracy of SNGD, even when given 40 times more training iterations. Given only 1,000 iterations, it would not reduce the loss or error much. Adam is omitted from sections 5.2-5.4 because PINNs often perform better with the use of second-order optimizers. However, we appreciate the reviewer’s concerns and will add ADAM to 5.2-5.4.\n\n- In our experiments, BFGS was more accurate than L-BFGS, so we omitted the latter. We would like to briefly note that although we don’t use L-BFGS in 5.4, we can compare to  “Challenges in Training PINNs: A Loss Landscape Perspective”, which uses the same benchmark problem. They report an error of $\\mathcal{O}(10^{-1})$ for a PINN using L-BFGS. In our manuscript, both BFGS and SNGD outperform this result.\n\n- We agree this should be included and will update the manuscript. We thank the reviewer for their suggestion. \n\n**Essential References Not Discussed:**\n\nWe thank the reviewer for these helpful suggested references and have added them to the related work section of the paper. \n\n**Other Strengths And Weaknesses:**\n\nWe agree that the presentation of the sketching-based approach should be improved. We are indeed using classical randomized SVD (in our case, an eigenvalue decomposition). We have clearly stated this in the manuscript and expanded our explanation of the sketching method. \n\n**Other Comments Or Suggestions:**\n\nWe thank the reviewer for pointing these out. We have fixed each error and enabled hyperlinks.", "claims_and_evidence": "Bottom of page 3: “the matrix $\\mathbf{G}$ should be rank-deficient”. The expression (6) for $\\mathbf{G}$ is an integral over the domain and boundary. In this case, why should $\\mathbf{G}$ still be rank-deficient? However, I believe that $\\mathbf{G}$ should be rank-deficient for the *discretized* PINN objective, which is what we actually optimize in practice.\n\nFigure 1: Does this exponential trend in the eigenvalues of $\\mathbf{G}$ hold for other problems? Moreover, which PDE does Figure 1 correspond to.\n\nTop of page 5: Why is the asymptotic upper bound of $10^{-16}$ valid? To make $\\lambda_{r + 1}$ small, $r$ might have to be large, which could make $C(p, r, e)$ large.\n\nI disagree with the authors' claim that CG-type methods can’t be used to solve the linear system. I agree that vanilla CG would struggle due to the lack of preconditioning. However, preconditioned conjugate gradient could greatly accelerate convergence within the linear system step. For example, a low-rank preconditioner could be incredibly helpful due to the spectral decay shown in Figure 1. Please see “Randomized Nystrom Preconditioning” (SIMAX 2023) for an example of such a preconditioner.\n\nWhat are the per-iteration complexity and memory requirements of SNGD? This is important to include so practitioners can understand whether Algorithm 1 will scale to their problem.\n\nIn what sense is the error in section 4.1 “near-optimal”? How do we know that a different sketching strategy can’t obtain a better error bound?", "methods_and_evaluation_criteria": "The algorithm itself makes sense for PINNs, since several previous works have shown that better optimization algorithms can lead to lower errors in solving PDEs.\n\nI’m concerned that the benchmarks in sections 5.1 to 5.3 are too simple. The analytical solutions for these benchmarks do not have high-frequency components, which suggests that they are relatively easy to solve. Perhaps the authors should run their experiments on other challenging settings, such as the wave equation from “On the eigenvector bias of Fourier feature networks: From regression to solving multi-scale PDEs with physics-informed neural networks” (Computer methods in applied mechanics and engineering 2021).", "theoretical_claims": "I don’t think Lemma 3.1 shows that the eigenvalues decay *exponentially*. Since the rate of decay is $j^{-4}$, the decay would actually be *polynomial*. For completeness, I would also recommend the authors provide a proof of Lemma 3.1 in the supplementary material, in addition to the proof sketch they provide in the main paper.", "experimental_designs_or_analyses": "Is the comparison in Table 2 fair? Table 2 makes it seem like Adam is slower than SNGD in the paper, but Adam is also run for 40 times as many iterations as SNGD.\n\nWhy is Adam omitted from the experiments in sections 5.2-5.4? I think the authors should also include Adam in these experiments in order to really demonstrate that their method is better than other optimizers for PINNs.\n\nI’m concerned by the authors not running L-BFGS in their experiments. L-BFGS is used more often than BFGS for training PINNs, since L-BFGS has a much lower memory requirement than BFGS. For example, “Characterizing possible failure modes in physics-informed neural networks” (NeurIPS 2021) and “Challenges in Training PINNs: A Loss Landscape Perspective” (ICML 2024) both use L-BFGS in their experiments, not BFGS. Without comparisons to L-BFGS, it is hard to say whether SNGD is actually a better method for PINN training.\n\nHow does single-pass SNGD perform on the transport equation? I would like to know how single-pass SNGD performs in this setting, since single-pass seems to outperform two-pass on the other benchmarks in the paper.", "relation_to_broader_scientific_literature": "To the best of my knowledge, the algorithm proposed in the paper is new. Given that PINNs are challenging to optimize, the algorithm proposed in this paper could be useful to practitioners. Despite my criticisms of the paper, I applaud the authors’ effort in developing a new optimizer for PINNs.", "essential_references_not_discussed": "The paper is missing two key references: “Sketch-Based Empirical Natural Gradient Methods for Deep Learning” (Journal of Scientific Computing 2022) and “Challenges in Training PINNs: A Loss Landscape Perspective” (ICML 2024). The “Sketch-Based” paper introduces SENG, which uses sketching-based techniques to approximate the natural gradient (although the sketching technique they use is different from the one in this submission). The “Challenges” provides an in-depth analysis of optimization challenges in PINNs and also proposes a new optimizer for PINN training, so it should also be discussed in the introduction and/or related work.", "other_strengths_and_weaknesses": "In general, I believe the presentation of the sketching-based approach could be improved. For example, the authors should say what kind of sketching-based technique they are using (I believe it is the randomized SVD, but please correct me if I am wrong).", "other_comments_or_suggestions": "I would recommend enabling hyperlinks to citations, figures, tables, etc. For example, when I click on a citation in the main text, it should take me to the appropriate location in the references.\n\nThere should be a period at the end of the second paragraph in the introduction.\n\nThere should be a space between “Descent” and “(SNGD)” at the top of Algorithm 1.\n\nIs there a typo in the single-pass step in Algorithm 1? I believe both $\\mathbf{Q}$ and $\\mathbf{M}$ have size $|\\theta| \\times (p + r)$, and this would mean that the matrix product $\\mathbf{Q} \\mathbf{M}$ is invalid.\n\nThe line search for $\\eta^\\star$ should use $\\theta_i$ instead of $\\theta$.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes to accelerate the energy natural gradient descent (ENGD) algorithm for PINNs through sketching via randomized SVD.
This approach is motivated by the spectral decay of the Gramian, which the authors study empirically and theoretically for a small toy problem.
Empirical results demonstrate that the resulting method, called sketchy natural gradient descent (SNGD), can train PINNs to higher accuracy faster compared to first-order methods and other second-order competitors like BFGS.

This work is borderline:

- The paper's idea to apply randomized linear algebra to the Gramian is straightforward and combines two known methods, yielding a simple algorithm.
  To the best of my knowledge, the paper is the first to apply this technique for training PINNs and the optimizer's performance compared to ENGD is promising.
  Some reviewers found the choice of using randomized SVD odd, as the Gramian is positive semi-definite and there exist sketching algorithms like the Nystrom approximation to exploit this structure.

- Some reviewers raised concerns that the presented experiments are rather small-scale.
  One reviewer requested additional large-scale experiments, suggesting concrete setups.
  The authors did not address this in their rebuttal, but I think this should have been possible given the proposed method's simplicity.

- Reading the paper, I found that there are some type-setting issues, e.g. some equations spill into the margins, and there are frequent spelling mistakes of ENGD (EGND).

Overall, I think the proposed method can serve as starting point for developing efficient second-order optimizers for PINNs using randomized linear algebra.
I am leaning towards recommending acceptance, but given its significant potential for further improvements, I would not mind if the paper was rejected.
Decision: Accept (poster)