Kronecker-Factored Approximate Curvature for Physics-Informed Neural Networks

Dear Reviewer aiiA,

thanks a lot for the time and effort you put into reviewing our work.

Regarding readability, we agree that layer indices can be omitted and will do so in the updated version. We will also think about other ways to make the notation lighter.

Weaknesses

[...] the loss curves only show the training loss and it is thus not clear if KFAC leads to more overfitting compared to other optimisers.

We respectfully disagree. We report relative L2 error rather than training loss, which avoids overfitting. The relative L2 error is defined as $\frac{\lVert u_\theta - u^\star\rVert_{L^2(\Omega)}}{\lVert u^\star\rVert_{L^2(\Omega)}}$, where $u^\star$ is the true solution of the PDE. In PINNs, the relative L2 error is usually considered as the relevant quantity and we estimate it on held-out points that are not used during training, similar to the test loss in supervised learning.

We apologize for not explicitly introducing the relative L2 error and have adjusted this in our revised version.

The loss curves do show that the KFAC approximation is often favorable, but only for a single run. It would be important to do multiple runs to average out the noise from different initializations and mini-batching, and also include error bars.

It is true that our plots show the performance of a single run. However, this run was obtained through extensive hyper-parameter tuning, and we found our results to be consistent when using a different search strategy (random or Bayesian). We think this already supports the robustness of our results, but understand your concern.

Since our computational ressources are limited, would you be satisfied if we presented an additional experiment with error bars for one of the sub-experiments?

I would have also liked to see some visualizations for some low-dimensional problems that show failure-cases of SGD and how KFAC fixes this. Or some other informative visualizations other than loss curves.

This is a great idea. We already have the code to do such visualizations because it was helpful to debug all optimizers. We will add such visualizations for the low-dimensional PDEs in the appendix.

Explicit questions

I noticed that the EMA update rate hyper-parameter is often super large, even up to 0.988 (for heat equation, appendix A9), which effectively looks only at the most recent batch. Do you have an explanation for why this is? What happens if you use a much smaller mini-batch size?

High EMA factors correspond to slowly forgetting pre-conditioners from previous iterations and are common in KFAC. For example, Martens and Grosse uses an EMA factor increasing to 0.95 (page 19), which is also the default value in KFAC's JAX implementation.

Is it correct that for linear PDEs, the approximation boils down to simply averaging over the additional forward passes corresponding to the additional derivatives on top of the average over the data points? So in comparison to the standard KFAC, you only have on additional average for these derivatives?

Yes, this is correct and very concise. We have added a clarifying sentence using your suggested phrasing.

For the non-linear PDEs, you had to make an additional approximation in C26 - C27 of the supplementary material. [...] Could you expand on this?

Yes, the GN-preconditioner corresponds to linearizing the PDE. The approximations (C26)-(C27) are analoguous to the linear case and do not require further approximations compared to the linear case.

The interpretation of $\Psi$ as a nonlinear output layer is correct. As $\Psi$ is not part of the convex loss, we have to linearize it for the GN-preconditioner. We group the Jacobian of $\Psi$ with the Hessian of the sample loss as it does not depend on the trainable parameters. It can also be grouped with the Jacobians $J_{n, \alpha}^{(l)}$, however, this results in the same approximation. We hope that this clarifies the question, but please do not hesitate to let us know whether we should elaborate on this any further.

It seems very odd that KFAC performs better if the random mini-batch are not sampled (new) at every step. Have you noticed any such thing for KFAC on standard problems (not PINN losses)? Are the matrices in KFAC inverted every iteration or also just every T-th iteration?

We are not aware of any findings in the literature that have observed something like this. To make our comparison fair, we allowed the number of iterations a mini-batch is used as a hyper-parameter on which we performed Bayesian optimization. We found that for all optimizers it was favorable to keep a mini-batch for significant number of iterations. Thus, we regard this finding as something specific to the PINN problem rather than to KFAC. Our methods update the Kronecker matrices and inverses at every step.

---

Thanks once more for your constructive comments. We hope our responses addressed them to your satisfaction. We remain attentive to your feedback!

Thank you for your response and for answering my questions.

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

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

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

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

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

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

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

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

Dear Reviewer YHUU,

we would like to thank you for the time and effort you put into reviewing our work. In the following, address some of the points that were rightfully raised by you.

How much additional memory is required [for our KFAC]? There is no per-step computational complexity provided

In our general response, we compare memory requirements and computation time of the proposed scheme. We find that the viewpoint of 'a larger net with shared weights' alias the forward Laplacian graph saves both memory and compute time when compared to the standard approach relying on nested backpropagation, as implemented in functorch. If there are any further questions, please do not hesitate to contact us.

No open-source implementation is provided

We will release the code after acceptance with the camera ready version.

K-FAC/K-FAC* do not reach the lowest error obtained by running ENGD [...]. Why is this the case?

As KFAC is designed as an approximation of ENGD we do not expect KFAC to provide results superior to ENGD. Note however that ENGD requires the solution of a system of linear equations in the parameter space of the network. Hence, while we can use it as baseline for a small network, it does not scale to larger nets. Our updated discussion of Figure 2 elaborates on this.

How would K-FAC/K-FAC* fair in more challenging settings [5, 6]?

We believe that PINNs should be employed for high-dimensional PDE problems where classical methods are intractable and are currently working on the log-Fokker-Planck equation which is a challenging nonlinear PDE. We aim to present first results in the discussion period and refined results for a camera-ready version.

The paper is missing more recent work on difficulties of training PINNs, such as [1, 2]

Thank you for the pointers to these references, which we have included into our introduction.

Notation/typos

I think notation is being reused too much in equations (5) and (6)

Thank you for this feedback. We are currently looking into ways to streamline and improve our notation for the camera-ready version.

I’d recommend defining the symbol as the Hadamard product before

Thanks for the pointer, we have added the definition.

Equation below line 168

Thanks for spotting this, indeed, we have been slightly inconsistent in carrying the subscript $N_\Omega$ for the number of interior integration points. We have adjusted this and have made our notation consistent.

the outer product decomposition in line 200 is missing a transpose

Thanks for catching this typo, which we have corrected.

Minor comments:

The sentence containing equation (3) is a bit confusing.

Thanks for the feedback. To our knowledge the matrix $J_\theta^\top u J_\theta u$ is usually called Gauß-Newton matrix, where $J$ is the Jacobian with respect to the model parameters. This is exactly the form of the matrix defined in (3) if $u = (u_1, \dots, u_n)^\top$. Please let us know, whether you have any further questions and we are willing to adjust our wording to reduce ambiguities.

I would recommend writing “K-FAC” and “L-BFGS”.

Thank you for the feedback, of course we are willing to follow the common nomenclature.

Questions

Line 64: Could the authors please elaborate on why the preconditioner not having a Laplacian term is beneficial?

We are deeply convinced that when optimizing a PINN, preconditioners with a PDE term, e.g., a Laplacian, are much more meaningful than preconditioners without a PDE term. In line 64 we mention a line of works that use KFAC for PINN, but only to approximate a preconditioner that does not include PDE terms. Hence, we regard these approaches as not satisfactory which motivated us to develop a KFAC method for preconditioners including PDE terms. We have slightly changed our wording in line 64 to improve clarity.

Line 131: Could the authors please elaborate on why the interior Gramian cannot be approximated with the existing K-FAC? Is this due to computational and/or autodiff limitations?

K-FAC has only been proposed when the loss and therefore the pre-conditioner only incorporates function evaluations. This excludes PDE terms in the pre-conditioner, which appear naturally when incorporating second-order information in PINNs. Hence, the reason existing KFAC approaches can't be used is not a computational limitation but quite simply the lack of a KFAC for problems including PDE terms.

Is using a general loss in equation (11) necessary?

We decided to work with a general loss as losses other than the least-square problem appear in practice. This is for example the case for variational approaches including the deep Ritz method and consistent PINNs which use Lp-norms. Of course, there is a trade-off between generality and accessibility. We decided to provide the general version as we already provide a didactic specific example with the Poisson equation and as we found the general form to be not much more complex. If you strongly disagree, we can also offer a simplified version in the main body and defer the fully general case to the appendix.

Does the proposed approach generalize to inverse problems involving PINNs?

We have not yet considered inverse problems. In general, we see no conceptual problem, however, one needs to take into account two (or more) neural networks when dealing with inverse problems. We leave this for future work.

Are there operators besides the Laplacian for which the weight sharing can be reduced?

A similar reduction can be obtained for linear scalar valued PDE operators, meaning whenever the PDE of consideration is a linear function of the partial derivatives of some order.

---

Thanks once more for your comments that have helped use improving our manuscript and we hope our responses could address your comments to your satisfaction. We remain attentive to your feedback!

Dear Reviewer TQnM,

we would like to thank you for the time and effort you put into reviewing our work. In the following, we want to address some of the points that you rightfully raised.

---

paper does not provide the complexity of the algorithm

does the method scale poorly with the input dimension? I would expect quadratic scaling with input dimension.

You are right that we currently do not discuss this; we address this in detail in our global rebuttal and will add a discussion to the main text.

The quick summary is that empirically we observe linear run time and memory scaling in the input dimension for computing the Laplacian either via functorch's nested first-order autodiff or the forward Laplacian framework. This linear scaling is intuitive from the forward Laplacian perspective, which performs input_dimension + 2 forward passes of the net (assuming the network's forward pass is independent of the input dimension, which is a good approximation for the MLP we used, but in general will depend on the architecture). In Landau notation, the scaling with the input dimension will depend on the choice of architecture.

Moreover, we observe some practical advantages of the forward Laplacian over the nested first-order autodiff approach:

• The constants in the scaling of the forward Laplacian are smaller than those in the functorch implementation. The forward Laplacian is roughly 2x faster. We think this is because its computation graph is less sequential. Also, we observe that the forward Laplacian uses roughly half the memory. These constants matter in practise.
• Backpropagation on the forward Laplacian graph (which we need to compute the gradient of the loss function) allows us to populate the K-FAC matrices without much re-computation.

Overall, it seems that the forward Laplacian is the current state-of-the-art method to compute input Laplacians, comes with advantages over nested first-order autodiff, and integrates nicely with our proposed K-FAC approximation.

---

it remains unclear how general is the approach, e.g. how to apply it to higher order PDEs or other neural network models.

Our approach is general:

• The proposed K-FAC approximation generalizes to arbitrary PDEs, both in the PINN formulation (i.e. least squares of the strong form) and for variational formulations (deep Ritz). We discuss this in Section 3.3, see also equation (14).
• The Taylor-mode AD needed to assemble the K-FAC matrices works for general neural network architectures. For instance, JAX offers a general purpose Taylor-mode implementation (jax.experimental.jet) onto which our proposed K-FAC framework could be built on.
• Like the traditional KFAC, our KFAC approximation is architecture-dependent. However, thanks to the recent formulation of KFAC for linear layers with weight sharing [2], a wide range of architectures is covered, e.g. fully-connected/attention/convolution layers.

---

Thanks again for helping us to improve our manuscript! We hope our responses address your comments. Please let us know if you have follow-up questions. We would be happy to discuss.

References

[2] Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., & Hennig, P. (NeurIPS 2023). Kronecker-factored approximate curvature for modern neural network architectures.

Thanks for your follow-up question!

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

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

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

Please let us know if you have any follow-up question.