ANaGRAM: A Natural Gradient Relative to Adapted Model for efficient PINNs learning

Thank you very much for your careful reading, insightful comments and constructive critics. We address hereafter the weaknesses you have identified and answer your questions.

Weaknesses:

The viewpoint of absorbing the differential operators into the model can also be understood as a Gauss-Newton method in function space compare to (Jnini et al, arXiv:2402.10680). Moreover, it leads to the standard Gauss-Newton algorithm in parameter space.

We thank you for this observation. We agree that the derivation provided in Jnini et al lead to similar algorithms but we would like to stress out some important differences:

• Our algorithm, while presented in its simplest form for the reasons explained in the general comment, is more general than Gauss-Newton. It is identical to Gauss-Newton only when neglecting two terms in equation (18) of the new version, one given by (66) in appendix which is usually null for linear operators ; (a fact highlighted in proposition 2 in the revision), and the other term introduced in equation (75) which is amenable to various levels of approximations, for instance using Nyström's method, the lowest one consisting to simply discard it.
• The choice of presenting only the simplified setting, coinciding with empirical Gauss-Newton was justified by the lack of competitive baseline in the literature. We thank the reviewer to point to our attention the paper by Jnini et al., but if we want to set a comparison with it, first we remark that the methodology is different, because in their case the solving of the least square problem is done either through conjugate gradient or direct least square regression with respect to the estimated gram matrix, while in our case we use SVD. On one hand the first option, to our best knowledge, does not allow a precise control of the cutoff factor, something that is nonetheless crucial for robustness of the algorithm, as illustrated by the aforementioned discrepancy between ANaGRAM and ENGD. Hence, we also suspect that convergence of conjugate gradient may fail in some cases due to the strongly bad conditioning, and that an additional preconditioner may be needed. On the other hand, direct least square regression with respect to the estimated gram matrix has a cubic cost with respect to parameters number, which is precisely what we aim to avoid. In practice, those are conjectures since the code seems not to be made available sofar according to page 6 of the paper.

The manuscript thus does not provide a novel algorithm, but a novel derivation of a known algorithm. While I like this idea and it is certainly a contribution, I think it is not significant enough for an exceptional venue such as ICLR.

We hope that the clarification of the theoretical framework, now introduced in section 3 and no longer in the appendix, will make a more meaningful contribution to you.

Questions

Can the authors add a reference for the computational complexity of the SVD for the rectangular solve?

We are not sure to understand this question since the computational complexity of (truncated) SVD is a classical textbook result of numerical linear algebra. One may for instance refer to Numerical linear algebra by LN Trefethen and D Bau, as suggested on the wikipedia page https://en.wikipedia.org/wiki/Singular_value_decomposition#Numerical_approach. Furthermore, recent works in numerical linear algebra bring new randomized algorithms that significantly improves the CPU time. We may for instance refer to the work of Professor Laura Grigori, in particular https://arxiv.org/abs/2302.07466 and https://arxiv.org/abs/2302.11474.

Why is the proposed method performing better than ENGD for the linear equations? I would expect similar performance as it is the same algorithm (up to numerical linear algebra).

The difference presumably comes from the fact that we took into account the cutoff factor of the SVD as an hyperparameter, which is for now adjusted manually on the logscale. This is something that was not taken into account in ENGD, explaining the discrepancy of performance for linear operators and underlying the importance of this hyper-parameter. The automation of it is under development in conjunction with the choice of collocation points.

It is very common to include damping into natural gradient methods (adding a scaled identity term to the square Gramian). Can this be realized in this setting?

Yes indeed. This can be done without further complication, just by replacing the term $\Delta^{\dagger}$ with $\frac{\Delta}{\Delta^2+\lambda}$, where $\lambda$ is the scaling factor.

Thank you very much for your careful reading, insightful comments and constructive critics. We address hereafter the weaknesses you have identified.

Weaknesses:

(1) The algorithm in the paper is not novel. The general natural GD method requires O(P^3) computational costs, due to the inversion of the Gram matrix of \hat{\phi} (G = \hat{\phi} \hat{\phi}^{T} is in size of P \times P). This paper applies SVD directly on \hat{\phi}. Moreover, the line 6 of the algorithm is exactly equivalent to the vanilla natural GD algorithm (rewriting the pseudo-inverse with the SVD). Therefore, I do not find the novelty of the algorithm.

(2) In Equations (18) and (19), it seems that the paper tries to explain the consistency by approximating the integral by finite sum of some observations. However, the results of Corollary 1 are relatively weak. The existence of certain points, as stated, follows naturally from standard regularity conditions (e.g., smoothness) of functions. The result is therefore not surprising and lacks substantial implications. It is better if the paper can alternatively prove that if we randomly pick P training samples as observations and with high probability, the difference between du and \hat{\phi} \hat{f} is upper bounded, e.g., O(1/P^{\alpha}), where \alpha is a constant depending on function’s conditions and dimensionality. Such findings would be more meaningful for machine learning.

(3) The ANaGRAM for PINNs closely resembles existing natural GD for training PINNs (https://arxiv.org/abs/2408.00573). For example, in Equation (31) of the paper, the pseudo-inverse of (J J^{T}) is equivalent to the SVD of \hat{\phi} in this paper.

If we understand well, we may reformulate critics (1), (2) and (3) by saying that our approach is just vanilla natural gradient with poor statistics. This is a legitimate interrogation which is probably reinforced by the simplified presentation of the algorithm that we gave which is possibly misleading, especially (former) equation (19), since it gives the false impression that we use a non controlled statistical criteria, while in fact our approach is a geometric one, namely the projection of the functional gradient into the empirical tangent space. This is now made clear in the new version in section 3, in particular with th.1 and eq. (18). We have tested the minimal version of Anagram by boldly neglecting two terms in (18), which already yield state of the art results. More precise algorithms can be considered based for instance on Nyström method in order to take into account these terms without significant additional cost, but we leave this for future developments. As a byproduct, this means that the quality of the batch is not given by a statistical criterion but by a geometrical one given by equation (18) in new version.

In the same spirit, former corollary 1 (now somewhat modified in proposition 1), which seems trivial, underlines this part, by showing that criterion of equation (18) can be saturated with P points. The exploitation of this criterion is under investigation as it represents a non trivial problem to optimize (18) while training the PINNs.

(4) The natural GD requires the full-batch training, which limits its adaptability to stochastic algorithms. We know that the SGD-like algorithms are widely used in deep learning due to their benefits, such as improving the generalization and bypassing spurious local minima. A more practical approach would be one that can adapt effectively to a stochastic setting.

As a follow up to previous answer, we believe that the introduction of the notion of empirical tangent space is a guide toward the adaption of natural gradient to a batched version. Various options are available, which are under investigation. We are aware that one of the main difficulties is to estimate effectively the dependence between different empirical tangent spaces. Nonetheless, we stress out that the gain with respect to classical SGD is already substantial, since SGD can be understood as projecting into $Span(NTK_\theta(x_i, \cdot))$ with respect to the push-forward metric of the Euclidian space into tangent space, while considering that points in batch are independent, which is very badly conditioned for PINNs as can be seen experimentally. In the revised version of appendices, we explain this fact with more precision, and provide empirical evidences of it.

Questions:

Please see the weakness part. One typo: Lines 126-127: ill conditioned

Thanks for pointing that out !

Thank you very much for the constructive remarks and kind comments ! We address here the weaknesses you have identified and answer here the questions you asked.

Weaknesses

There are many typos 126 il conditioned -> ill condition 186 hapens -> happens 201: writes -> write 274 quadraic.... The flow of the paper may be better by revisiting some parts. For example, at the end of the section 2, authors illustrate the relationship between the natural gradient and the linear regression problem. However, the section 3.1 has the title indicating the same thing. Maybe the end of section 2 can be moved to section 3.1, otherwise, I don't see the reason why section 3.1 has the current titles.

Thank you very much for all these comments. We have revised the manuscript accordingly. In particular Section 3 has been entirely revised.

The experiments are not enough. Usually people also include Burgers, Schrodinger, etc.

We thought that our choices represented well enough the different study cases (linear, non-linear, higher dimension ...). But we can definitely also include those experiments.

The paper's main theoretical result is the equivalence of Green's function and natural gradient. It would be helpful to include some toy examples to justify this finding.

Thank you for this suggestion, we can indeed illustrate this on a set of Fourier features, as well as with figures. We would add this int the revised appendices. That being said, we would like to emphasize that in our opinion there is a second as important theoretical result, which is the introduction of the notion of empirical tangent space and empirical natural gradient, presented in appendix section C, and that we now highlight in section 3 of the revision.

Questions:

It's interesting that ANaGRAM usually has the largest variance among all methods, is there a reason for that? I was thinking it should have small variance since it's more well conditioned.

This is only an "optical illusion" due to the logscale. You can have a look at section A.3 in supplementary material where a detailed overview of the numerical results is given.

For some experiments, it seems the difference between different methods is more stark for L2 loss compared to test Loss. Any explanation for this?

Thanks for this observation. We suspect that this can been explained by the fact that the "residual norm", which is represented by the test loss and the L2 norm have different spectral biases. There is a line of works studying mathematically these kind of effect for which we are not too familiar but which might be relevant to explain that. We refer for instance to the dissertations of Tim De Ryck (https://www.research-collection.ethz.ch/bitstream/handle/20.500.11850/674112/dissertation_deryck.pdf?sequence=1) or Roberto Molinaro (https://www.research-collection.ethz.ch/bitstream/handle/20.500.11850/646749/Thesis%2813%29.pdf?sequence=1).

Could you explain more about the relationship between NTK and equation 19 (line 223)?

The relationship is explained in appendix C. Nevertheless, since this relation is indeed very important and that equation (19) (in the original version) seems to be misleading, in the sense that it hides the geometric interpretation under the NTK relationship, we have completely revised this part to discuss the core of Anagram which is now presented in section 3 and encapsulated in theorem 1. and its consequence like e.g. eq. (18) in the new version.

We would like to thank you for your interest and the helpful remarks and comments ! We address here the weaknesses you have identified and answer the questions you asked.

Weaknesses:

1. The current implementation of ANaGRAM requires manual tuning of hyperparameters, such as the cutoff factor for SVD.
2. The choice of collocation points (batch points) is limited to heuristics in this study. Developing an adaptive method for selecting these points could further improve ANaGRAM's performance.
3. The results indicate that the efficacy of ANaGRAM may diminish for nonlinear differential operators, where the method's equivalence with traditional approaches is less certain.

The observed diminution of efficiency for nonlinear operators can be thoroughly explained by the theoretical framework we developed, namely the error stated in equation (66) of the appendix, now the second term stated in theorem 1, that is null for linear operators as stated by Proposition 2 in revised version. Note that a second approximation is made, introduced in equation (75), now the first term stated in Theorem 1.

Questions:

1. Implementing an adaptive selection strategy for the SVD cutoff factor based on the eigenvalue spectrum could make the method easier to use.
2. Developing an adaptive method for selecting collocation points based on the geometry of the domain could improve convergence rates.

We do agree on both remarks. Nevertheless, we do think that this is a different line of work, that we are currently carrying on. As already mentioned in the general comment, our main objective was to put on firm ground the natural gradient in PINNs context, so as to have a sound basis on which to develop future extensions (e.g. collocation points, hyperparameters tuning). This paper constitutes for us the first step on which we can elaborate, as it gives us in particular a good criteria for collocation points optimization, namely equation (18) in the new version.

We would like to thank you very much for the positive and supportive appreciation of our work, as well as your helpful remarks and comments ! We particularly appreciate the acknowledgment on mathematical work ! We address here the weaknesses you have identified and answer here the questions you asked.

Weaknesses

The performance of ANaGRAM is quite similar to E-NGD for Heat eqn and 5D Laplace eqn. ANaGRAM is faster by a tiny margin. I did not find a discussion of this in the paper. Can the authors comment on this? In particular, what would be the expected advantage of ANaGRAM over E-NGS in words?

The equivalence between the two algorithms in the case of linear operators is the subject of section E in the appendix. The difference is that ANaGRAM has a better complexity when the number of parameters is significantly larger than the batch size, which is not the case in these experiments. We will try to illustrate that on the High-dimensional (100) Laplace example. Also in the non-linear case, the mathematical equivalence between both algorithms does not hold anymore as one may see in the Allen-Cahn example.

Questions

I find the empirical investigations sufficient for this paper. It might be even better to have a high dimensional example. For example, would ANaGRAM still perform better than E-NGD when tested on 100 D Laplace eqn?

We did not considered such a high-dimensional equation, as it seems to us a bit far from numerical analysis concerns, for which PINNs are meant. Nevertheless, we will try to investigate this regime.