Near-optimal Sketchy Natural Gradients for Physics-Informed Neural Networks

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.

Claims And Evidence:

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

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

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

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

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

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

Methods And Evaluation Criteria:

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

Theoretical Claims:

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

Experimental Designs Or Analyses:

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

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

• We agree this should be included and will update the manuscript. We thank the reviewer for their suggestion.

Essential References Not Discussed:

We thank the reviewer for these helpful suggested references and have added them to the related work section of the paper.

Other Strengths And Weaknesses:

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

Other Comments Or Suggestions:

We thank the reviewer for pointing these out. We have fixed each error and enabled hyperlinks.

We thank the reviewer for their helpful and detailed suggestions and comments and the time and effort they put into reviewing our manuscript.

Experimental Designs Or Analyses

We will add comparisons with ADAM to 5.2-5.4.

At the end of section 5, we specify how $p$ is chosen and set the tolerance to 1E-13 across all experiments.

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

Essential References Not Discussed

Regarding “Streaming …” please see the comment (2) under Weaknesses below.

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

Weaknesses

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

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

Concretely, 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) \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. We are not aware of a way to express $\mathbf{G}\_{i+1}$ as $\mathbf{G}\_{i+1} = \eta\mathbf{G}_i+ \nu \mathbf{H}$.

Other Comments Or Suggestions We thank the reviewer and have updated the manuscript to fix each error.

Questions For Authors

1. We have fixed a typo in Algorithm 1: $Q \gets Q[:,1:(p+r)]$, so $T = Q^{T} G Q$.
2. 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}$.
3. 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.
4. We have added some discussion in the last paragraph of section 4. 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. 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.

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.

Relation To Broader Scientific Literature:

We thank the reviewer for the suggestion and have included a brief discussion of this reference in the manuscript in the “related work” section.

Weaknesses:

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

Questions For Authors:

1. Damping can be included in the solution approach, and we thank the reviewer for the suggestion. We will add it to future work.
2. 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.
3. • 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.
   • 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.
4. 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.

We thank the reviewer for taking the time to carefully review our manuscript and appreciate the positive review of our work.

Questions For Authors:

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

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

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