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

Authors: Nilo Schwencke, Cyril Furtlehner
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: PINNs, SciML, PDEs, Natural Gradient, Neural Tangent Kernel

Abstract:
In the recent years, Physics Informed Neural Networks (PINNs) have received strong interest as a method to solve PDE driven systems, in particular for data assimilation purpose. This method is still in its infancy, with many shortcomings and failures that remain not properly understood.
In this paper we propose a natural gradient approach to PINNs which contributes to speed-up and improve the accuracy of the training.
Based on an  in depth analysis of the differential geometric structures of the problem, we come up with two distinct contributions:
(i) a new natural gradient algorithm that scales as $\min(P^2S, S^2P)$, where $P$ is the number of parameters, and $S$ the batch size;
(ii) a mathematically principled reformulation of the PINNs problem that allows the extension of natural gradient to it, with proved connections to Green's function theory.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 3, 3

Reviews and discussion:
<thread>
creation date:
1730386584646

thread_type:
official_review

<thread_object>
creator: reviewer_ogx7

<content>
summary:
The manuscript proposes a natural gradient method related to the ENGD optimizer proposed by (Müller and Zeinhofer). The main contributions are:

- In the context of PINNs, the authors view the ansatz/parametric model as a neural network followed by the application of the PDE and boundary operator. This means that the model is defined on some parameter space $\Theta$, taking values in the product space $L^2(\Omega) \times L^2(\partial\Omega)$ via the intermediate space $\mathcal H$ which is related to the PDE and boundary operators of the specific problem. This mathematical trick of changing what is the model and what is the loss function allows to view PINNs as a standard regression problem. This viewpoint motivates the authors to use an $L^2(\Omega)\times L^2(\partial\Omega)$ based natural gradient method, which agrees with ENGD for linear PDEs but performs better than ENGD for nonlinear PDEs

- The second contribution is on the numerical linear algebra level. Instead of solving the normal equations of the regression, the authors propose to use an SVD directly for the Jacobian of the model. Here, when I write Jacobian I mean the discretized differential of the parametrization, which is denoted by $\hat \phi_{\theta(t)}$ in the manuscript. This gives an improvement in computational complexity from $\mathcal O(P^3)$ for solving the normal equations, to $\mathcal O (\min(N^2P, P^2N) )$, as the Jacobian is of shape $(N, P)$, where $P$ is the dimension of the parameter space $\Theta$ and $N$ is the number of sampling points.

strengths:
- The viewpoint outlined above leads to a performant method, that works better than ENGD for nonlinear PDEs.

- The improvement in complexity for the linear solve is helpful.

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

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

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

- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730479612811

thread_type:
official_review

<thread_object>
creator: reviewer_inht

<content>
summary:
This paper studies the natural gradient descent, which is theoretically and empirically believed to improve the convergence of training PINNs. Theoretical analysis on the natural gradient and training PINNs bridges connections on Green function theory.

strengths:
(1) The natural gradient descent approach is theoretically sound, offering a promising solution for improving the convergence of regression problems and PINNs. Notably, while gradient descent may struggle with PINNs involving nonlinear PDEs, natural gradient descent effectively addresses this limitation.


(2) The paper's presentation is clear and well-structured.

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. 


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

questions:
Please see the weakness part.

One typo: Lines 126-127: ill conditioned
</content>
</thread_object>

</thread>
<thread>
creation date:
1730514088251

thread_type:
official_review

<thread_object>
creator: reviewer_uhpu

<content>
summary:
This paper proposed a natural gradient based optimizer called ANaGRAM for PINNs. The key algorithmic contribution of the paper includes a novel approach to discretize the natural gradient by connecting the natural gradient with a linear regression in L2 space, which is then combined with the fact that PINNs can be regarded as a least square regression to adapt the method for PINNs. A major theoretical contribution of this is showing that the natural gradient for PINNs is actually Green’s function under some conditions. 

In summary, this paper proposed a novel optimization approach which can be used to help improve the convergence of the usually ill-condition problems in PINNs.

strengths:
This paper has an interesting and supportive theoretical results for their results.

weaknesses:
1. There are many typos 126 il conditioned -> ill condition 186 hapens -> happens 201: writes -> write 274 quadraic....
2. 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.
3. The experiments are not enough. Usually people also include Burgers, Schrodinge, etc.
4. 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.

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

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

3. Could you explain more about the relationship between NTK and equation 19 (line 223)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731101396881

thread_type:
official_review

<thread_object>
creator: reviewer_eamy

<content>
summary:
The paper proposes a new algorithm -- ANaGRAM-- that is well suited for learning PDEs with boundary conditions. The algorithm makes use of a clever manipulation of natural gradient of Amari & Douglas. The natural gradient, PINN and some elements from differential geometry are introduced. The proposed algorithm can be interpreted as follows: compute the integration of the euclidean gradient using by "weighting" with the appropriate Green's function and use pseudo-inverse of the differential as a pre-conditioner (Theorem 1). Experiements show that ANaGRAM learns faster than Adam, GD, E-NGD across well-known low-dimensional PDEs.

strengths:
The paper is very well written. The proposed method is elegant and well motivated from the theory perspective. Empirical results are well-done and show good evidence for the proposed ANaGRAM.

weaknesses:
I quite like the paper. Even though I am not an expert of PINNs at all, I am guessing that this paper may be exciting for that community. I found the math well written, which may be not very common for an ICLR submission. 

Typos:
- Eq 2: $h_i \to h_p $
- line 147: "his" comp. cost $ \to $ "its" comp. cost
- line 199: maybe add what $\Delta_{\theta_r}$ means here (singular values) 

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?

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731128798747

thread_type:
official_review

<thread_object>
creator: reviewer_pscc

<content>
summary:
The paper "ANaGRAM: A Natural Gradient Relative to Adapted Model for Efficient PINNs Learning" introduces an innovative natural gradient optimization algorithm tailored for Physics-Informed Neural Networks (PINNs). Named ANaGRAM, this approach leverages differential geometric perspectives to optimize PINNs, aiming to enhance computational efficiency and accuracy in solving Partial Differential Equations (PDEs) in a machine learning framework.

strengths:
1.The paper's use of differential geometry to formulate a natural gradient specifically for PINNs counts as a theoretical contribution. By aligning gradient updates with the functional space geometry, ANaGRAM offers a alternative to gradient-based methods, potentially addressing some of the known limitations in PINN optimization.

2 NaGRAM's computational complexity, which scales favorably with parameter and batch sizes, is an improvement over traditional natural gradient methods.

3. Experimental results indicate that ANaGRAM consistently achieves lower error rates compared to baseline methods like E-NGD and standard optimizers (Adam, L-BFGS).

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732151522478

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General comment

comment:
We thank the reviewers for their careful reading of the paper and their insightful comments and constructive critics. The main ones can be summarized as:
1. the method does not seem to bring any novelty, as it boils down to Gauss-Newton
2. the method is suboptimal in many respects

## Concerning the first line of critics
We disagree with the reviewers because in our view Anagram is much broader than that, as it offers a different perspective on the natural gradient, by shifting the preconditioning from the parameter space to the "sample space" which is interesting in the context of PINNs since the batch points can be chosen arbitrarily. This leads to the notion of empirical tangent space and empirical natural gradient the latter being geometrically interpreted as the projection of the functional gradient on the empirical tangent space. This projection can be made more or less approximate, and it is true that the version we presented corresponds to the crudest approximation, which indeed happens to be equivalent to empirical Gauss-Newton adapted to the PINNs context (we thank reviewer ogx731 for pointing this to us). We made this choice for two reasons, first because we were already obtaining  empirical results much better than existing baselines, on precision and speed either on one or on both criteria by a large margin; secondly we wanted to insure  a clear presentation in the core paper for pedagogical purposes, which was something that seemed to be appreciated by some of the reviewers, leaving unfortunately the discussion of some key points in appendix which led seemingly to some misunderstanding. Indeed the drawback of this simplified presentation is that it possibly gave the false impression of non original or trivial results.
Summarizing the contributions we have:
* a theoretical framework which introduce the notions of empirical tangent space and empirical natural gradient leading to a family of algorithms (Anagram) which depends on the way the projection of the functional gradient on the empirical tangent space is approximated;
* a key relation showing that our formulation of  Natural Gradient [resp. empirical Natural Gradient] for PINNs coincides with the operators Green function restricted to the tangent [resp. empirical tangent] space, which is also to our best knowledge completely new. This demonstrates  the coherence of the framework and that in some sense it is an optimal order 1 algorithm for PINNs. % will be difficult to define a better natural gradient for PINNs.
* a basic but efficient implementation of the simplest  instantiation  of Anagram, which indeed can be seen as a combination of Gauss-Newton with SVD adapted to PINNs, with good scaling properties, showing robustness and superior empirical results to existing baseline and this by a large margin either on speed or precision or both. We also provide a formal %bound
estimate on the projection error in this case; 
* a new, simple and principled optimization criteria for the collocation point problem, which is a direct byproduct of our theoretical framework. 

## Concerning the second line of critics
We indeed recognize that many aspects of the method can be improved (choice of the points, heuristic for the cut-off, optimized svd ...). 
Actually our focus was to provide a theoretical principled picture addressing the spectral bias, while stressing out the benefits in terms of complexity. Compared to previous natural gradient methods for PINNs (in particular in Jini et al : http://arxiv.org/abs/2402.10680 and Müller and Zeinhofer : https://proceedings.mlr.press/v202/muller23b.html), our method comes with a generic way to lower the cost of natural gradient and an explanation of its optimality through the equivalence with Green's function. What we see experimentally is that we already reach better than state of the art performances without inserting side sophistication in the algorithm, which empirically validate our theory. 

We provided a modified version of the paper aiming at highlighting more our contributions listed above and at dissipating some misunderstanding. The control of the empirical gradient projection leading to the Anagram framework is now specified by theorem 1 and the cost function for the batch points which is obtained as a byproduct of the theory is introduced in eq (18). The modification of appendices will be provided soon, taking into account additional developments asked by the reviewers.
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) **Summary of Claims and Findings**  
The paper proposes ANaGRAM, a novel optimization algorithm based on a natural gradient approach tailored to Physics-Informed Neural Networks (PINNs). The primary contributions are:
1. A theoretical framework that connects the natural gradient for PINNs with Green’s function theory.
2. A new formulation of the natural gradient leveraging the empirical tangent space, enabling scaling relative to parameter and batch sizes.
3. Empirical results demonstrating that ANaGRAM improves the speed and accuracy of PINN training compared to baselines such as Adam, L-BFGS, and E-NGD.

(b) **Strengths**  
1. **Theoretical Contribution:** The introduction of the empirical tangent space and its relationship to natural gradients and Green’s function theory provides a solid foundation for the approach.
2. **Algorithmic Design:** The use of SVD for preconditioning improves computational efficiency while maintaining theoretical rigor.
3. **Empirical Validation:** Experimental results consistently show superior performance over baseline methods in terms of error rates and convergence speed.
4. **Clarity:** The paper is well-structured, and the theoretical arguments are well-articulated.

(c) **Weaknesses**  
1. **Limited Numerical Exploration:** Although the theoretical framework is robust, the empirical work could include more diverse PDEs and higher-dimensional problems to strengthen the claims.
2. **Manual Hyperparameter Tuning:** Key parameters such as the SVD cut-off factor are manually tuned, which could hinder usability in practice.
3. **Criticisms on Novelty:** Some reviewers argue that the method resembles existing natural gradient approaches, albeit with nuanced improvements. This limits the perceived originality for an exceptional venue like ICLR.
4. **Adaptability to Stochastic Settings:** The method currently requires full-batch training, which is less adaptable than stochastic approaches like SGD.

(d) **Decision Justification**  
The paper provides significant theoretical insights and empirical evidence for improving PINN training via natural gradient methods. While concerns regarding novelty and numerical breadth are valid, the methodological and theoretical advancements make this work a meaningful contribution to the field. The authors’ rebuttals adequately addressed the reviewers' main concerns, emphasizing the geometric interpretation of the empirical tangent space and the innovative use of Green's function theory. The decision to accept is based on:
- The strong theoretical contributions.
- Empirical performance demonstrating improvements over state-of-the-art methods.
- Potential impact on the PINN and scientific machine learning communities.