Paper: Spectral-Bias and Kernel-Task Alignment in Physically Informed Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The clarity of the paper needs to be improved. As a theoretical paper, one cannot tell the story by throwing out equations without explaining its assumptions. The two main results (i.e., NIE and spectral bias) need to be formalized into theorems.

Author response:
We thank the referee for those comments. We agree that too many derivation details were carried through to the appendix. As a result, we now added a new derivation outline section (section 2 in the revised version) which goes through the main approximations involved in obtaining this result for both linear and non-linear differential operators.

Reviewer comment (Evaluation):
The paper needs more discussion about why the main results are important. In `section 2.1` and `section 2.3`, a lot of focus has been put on discussing how the main results are derived. However, an intuitive explanation of why they are useful is missing.

Author response:
In this work, we derive theoretical results that give the dataset-average prediction of a PINN. It is a fundamental building block of what we hope would be a fruitful theoretical investigation into spectral bias. Similar works in the context of supervised learning (https://www.nature.com/articles/s41467-021-23103-1) have received much and justified attention. 

That being said, by quantifying spectral bias, we can specify what PDE-solutions are best aligned with the DNN's architecture. This is the context of Fig. 2. showing that (a) performance correlates with how well the target lays in the dominant eigenmodes of the kernel. As shown there, in panel (a) when the cumulative spectral distribution function increases slowly, meaning that the target is spread over many kernel eigenmodes--- performance deteriorates. In panel (b) we demonstrate that more data allows for the learning of lower kernel eigenmodes as reflected in the cumulative spectral function of the learned function. 

A discussion in this spirit has been added to Sec. 4.

Reviewer comment (Theory):
As far as I understand, the spectral bias in this paper is based on the assumption that the PINN is infinitely wide. In practice, of course, every NN has a finite width. It is not clear how much things will break down when we look at a finite-width NN.

Author response:
The referee is asking a very good question which has been studied in depth in the context of supervised learning (see for instance https://proceedings.neurips.cc/paper/2020/file/ad086f59924fffe0773f8d0ca22ea712-Paper.pdf). For fully connected networks, such as those commonly used in PINN, the performance of infinite-width GP limits is often on par or better than the actual finite neural network. This seems to be related to the lack of feature learning in such networks. Turning to CNN, finite networks do excel over their infinite-width counterparts, but GPs still seem to be a good approximation. In particular, the benefit of clever architecture choices is reflected in the infinite-width GP limit (see for instance https://arxiv.org/abs/1810.05148 for the beneficial effect of pooling on both CNN and their GP limits)

Reviewer comment (Writing):
Since the main text is 7 pages right now, I recommend the authors add more introduction to the background of this work. The latter part of `section 1` goes a bit too fast and cannot be easily comprehended without a solid familiarity with theory of PINNs.

Author response:
We thank the referee for pointing out these issues. The current version now contains far more introductory material. In particular, a detailed explanation of the PINN setting and the training dynamics we use, an introduction to our use of replicas, and dataset averaging techniques. Further demonstrations of the NIE for non-linear operators and the variational approach we take.

Reviewer comment (Theory):
In your spectral bias statement, can you make any general statements about what the eigenvalues and eigenvectors will look like, or is it a very case-specific thing?

Author response:
The referee is asking a very good question. On the one hand, spectral bias reflects both the DNN's architecture and the dataset measure. In supervised learning, architecture and datasets play a major role in DNN performance. We thus expect spectral bias to be a rather rich quantity that adapts to the data and architecture in intricate ways. In these circumstances, one must resort to numerical method to detect and interpret spectral bias (specifically diagonalizing the kernel on a large and finite dataset, finding an ansatz based on a linear combination of outputs-per-neuron of the DNN which matches the eigenvector on the finite dataset but allows extrapolating to arbitrary points).  

Luckily PINNs are much simpler in this context. First of all, dataset measures, at least in vanilla settings, are uniform and low dimensional. Hence, one can directly visualize the eigenfunctions of the kernel (unlike in say CIFAR-10 where these would be non-linear functions in dimension 3072) and focus on the architecture's effect on spectral bias. For FCNs and data domains (i.e. rectangular boundaries or no boundaries), we generally expect these modes to correspond roughly to Fourier modes and the difference between architectures would lay mainly in their different spectral decay. So, for instance, trying to use it for equations with multiscale phenomena such as in turbulence,  one would need to design FCN with a much slower spectral decay than in our FCN experiments. This can potentially be done by taking very deep FCNs and tuning these to the "edge of criticality". Still, an experiential demonstration of this is beyond the scope of the current work.

Reviewer comment (Experiments):
experiments are too simple.

Author response:
Perhaps this is related to the issue with the loss function. Our theory paper derives, for the first time (as far as we know), the expected generalization error of a PINN network averaged over all possible dataset draws. Similar theoretical works, in the realm of supervised learning, have received much attention [e.g. https://www.nature.com/articles/s41467-021-23103-1]. We thus strongly feel that a theoretical gap has been filled here. The experiments we carry out are mainly meant to demonstrate the theory at its most basic level.

Reviewer comment (Novelty):
lack of reference, such as  
Neural-network-based approximations for solving partial differential equations. communications in Numerical Methods in Engineering, 1994  
Training behavior of deep neural network in frequency domain.  
On the spectral bias of deep neural networks.  
The convergence rate of neural networks for learned functions of different frequencies.  
A fine-grained spectral perspective on neural networks.  
On the exact computation of linear frequency principle dynamics and its generalization.

Author response:
We thank the referee for pointing out these references. Due to lack of time and the desire to reference these properly, we will add those to the next version.

Reviewer comment (Novelty):
what is the advantage of the proposed form?

Author response:
It is not clear to us what is the comparison here. An advantage over what? currently, there is no similar analytical equation governing the PINN data-averaged output. That being said, several parts of the manuscript including Fig. 2. are dedicated to showing how the spectral bias encapsulated in this equation affects learning.

Reviewer comment (Presentation):
Fig. 2 is hard to be understood. I also do not know how to use spectral analysis to understand the different results.

Author response:
At the basic level, the source term in the equation should be spanned by high eigenfunctions of the $L \hat{K} L^{\dagger}$ kernel operator. In various simple cases (e.g. linear equation spatially homogenous apart from a source term, and set in a simple domain), these eigenfunctions are roughly Fourier modes and the statement simplifies to saying that simple PINNs work better when the source has only low-frequency modes. Even then, there is the question of how fast do the eigenvalues associated with higher Fourier modes decay. An exponential decay with wave number, as we observe for simple networks, would make it very difficult to capture multiscale phenomena such as turbulence or shock waves.  

Next considering richer domains, inhomogeneous operators, etc... our results guide the practitioner to what is the right analog of high Fourier modes. The practitioner can then try and ``spectrally engineer" the kernel such that the source term has larger support on the (analogous-to) high Fourier modes. Such support is captured by our cumulative spectral function graphs ($A_k$) where a faster-increasing function means better support and better expected performance. 

Figure 2 is a demonstration of this spectral dependence of learning. The function being learned by the PINN is predominately made of high spectral kernel eigenvalues and lower modes are slowly learned as more and more datapoints are given. 

While the rough notion of spectral bias is expected in PINN, our work places it on a formal and quantitative theoretical footing. 

We note by passing (and can provide an appendix on this) that one can readily visualize these eigenfunctions by: 1. Obtain the kernel between $n$ data points. Diagonalize the matrix to obtain eigenvectors. Find a linear combination of network's outputs per neuron which coincides with the former eigenvector on the discrete set of datapoints. Note that, provided $n$ is larger than the width of the network, this is a well-posed problem. Use the resulting linear combination to approximation the kernel eigenfunction. Plot it on the domain. 

Following the various additions to the text, we hope the referee will give our paper a second read.

Reviewer comment (Presentation):
The main weakness of the paper is not clear enough presentation, making the accurate reading of the paper (e.g. by trying to verify theoretical results) quite difficult. Given the almost 2 unspent pages of the main paper and quite the small size of the current appendix, there is a lot of room for giving the necessary details.

Author response:
We agree with the referee that too many derivation details were carried through to the appendix. As a result, we now added a new derivation outline section (section 2 in the revised version) which goes through the main approximations involved in obtaining this result for both linear and non-linear differential operators.

Reviewer comment (Writing):
The setting is not described clearly and in full detail. For example, the authors mention that posterior (3,4) describes infinitely wide NN trained by Gradient Descent with white noise the gradients to the gradients weight decay. Apparently, this setting follows [Naveh2021], but this is not mentioned explicitly by the authors, and the respective reference is mentioned together with classical NTK work [Jacot2018]. Since the setting with added white noise to the gradient is much less common compared to that of [Jacot2018], it may take a significant time to figure out what the authors actually mean. Moreover, as mentioned by [Naveh2021], the considered setting is equivalent to Bayesian training with NNGP kernel, adding more options for confusion.

Author response:
Following the referee's advice, we now clearly state what training protocol we follow and how it matches Bayesian inference with Gaussian-noisy observation. We comment that this training procedure is close in spirit to NTK. Both approaches take infinite width, vanishing learning rates, and MSE loss. Our version adds noise to the gradients and weight decay to make closer contact with Bayesian sampling as $t \rightarrow \infty$. For any finite $t$, taking the noise to zero, the two training protocols coincide. In practice, our Bayesian posterior, also referred to sometimes as Bayesian posterior with the NNGP kernel, is on par with NTK in terms of performance as shown in various works [see for instance https://proceedings.neurips.cc/paper/2020/hash/ad086f59924fffe0773f8d0ca22ea712-Abstract.html].

Reviewer comment (Theory):
While it is expected that the result of [Naveh2021] for standard fully-connected networks carries over to the physically-informed networks and results in the posterior in eq. (3,4), the clarity of presentation would be improved if the authors would give the reference for the respective statement (if it was obtained previously) or formulate and prove a proposition (even if it turns out to be quite simple).

Author response:
The revised setting section now hopefully clarifies this point better. See specifically Eq. (3).

Reviewer comment (Writing):
Derivation of eq. (5) uses the replica trick to evaluate the average of possible choices of the training dataset $\mathcal{D}_n$. I would expect that the replica trick is a technique that a typical member of ICLR community does not know by default, so at least a reference or a short description of the technique is needed to make the derivation accessible to a broader audience.

Author response:
The new "outline of derivation" section (Sec. 2. 2nd paragraph) now hopefully discusses this more clearly.

Reviewer comment (Theory):
After decoupling replica modes, the authors continue to work with the effective action (33) of a single mode. This raises questions: (i) whether the replica limit $p\to0$ was taken implicitly at some point and (ii) why this effective action describes the (data averaged) posterior distribution of $f$ we are interested in.

Author response:
The referee is right to point out this step which the text unfortunately glanced over. The revised version contains an extra equation explaining this (see discussion below equation 37). Specifically, once replicas are decoupled, the data-averaged replicated partition function ($\langle Z^p \rangle_{data}$) takes the form, $\langle Z \rangle_{data}^p$ hence yielding $\log(Z) =_{\lim p \righarrow 0} (1-\lange Z^p \rangle_{data})/p=\log(\langle Z\rangle_{data})$. Thus we can use the single replica partition as the posterior partition function.

Reviewer comment (Theory):
There is a lack of clarity on how NIE equation (5) can be rewritten in the form (46) without integration by parts, and more details are needed.

Author response:
This is due to our $L^{\dagger}$ notation, where we state that it means $L$ is acting to the left ($fL^{\dagger}...=(Lf)...$. Thus, we haven't really carried the derivative to act on the right side and formally avoided the integration by parts. For instance if $L=\partial_x$ then our notation means $\int_0^1 dx f L^{\dagger} g \equiv \int_0^1 dx L[f] g$, then following an integration by parts it is equal to $-\int_0^1 dx F L[g]+ fg|_0^1$

Reviewer comment (Theory):
There are a number of assumptions and approximations performed in the paper, and more careful commenting and/or validation of these assumptions would make the claims of the paper more verifiable.

Author response:
The revised manuscript now contains a simple yet typical example of how we can tackle non-linear operators (App. Sec. B3). The new outline-of-derivation section should also provide more clarity on this matter.

Reviewer comment (Theory):
The statement that the proposed approach can be extended to the nonlinear operators $L$ is supported only by a single line at the end of Appendix B.2, which hints at a potential derivation strategy for nonlinear operators $L$. If not described in more detail (e.g. by providing execution of the derivation strategy suggested by the authors), this claim seems suitable only for the discussion section (as an example of a direction for future work).

Author response:
The revised manuscript now contains a simple yet typical example of how we can tackle non-linear operators (App. Sec. B3). The new outline-of-derivation section should also provide more clarity on this matter.

Reviewer comment (Presentation):
Regarding the plots, there is a small issue of visibility of axes and colormap tick values on the inset heatmap plots from Figure 2, which are required to support the benefit of stronger concentration of cumulative spectral functions. At the moment, the only way to see them is with extremely high zoom and good eyesight.

Author response:
Would this manuscript be accepted, we promise to address these graphical comments which we cannot do at the moment due to technical constraints. We are very thankful to the referee for his/her important insights and careful reading of the text.

Reviewer comment (Writing):
Minor typo — tilde is missing in the right-hand side of the reparametrization of variation from $h$ to $\widetilde{h}$ (eq. after eq. (35)).

Author response:
Many thanks again. Both have been corrected.