Infinite Neural Operators: Gaussian processes on functions

Thanks for your thoughtful feedback and contribution towards sharpening our definitions and fleshing out the assumptions behind our theoretical results . Below, we provide clarifications regarding the content of Section 3, as well as our experiments in Section 5.2.

We hope our answers sufficiently address your concerns and elevate your appraisal of our work. Otherwise, we would be happy to engage further.

What is meant by the covariance function associated with a particular operator?

We apologize, the statement of Lemma 3.1 in the main text is missing some details that are only present in Appendix B.1. This will be corrected and improved in the main text. In summary, $B\_1$ is a Hilbert space-valued stochastic process, which is not necessarily Gaussian, while $B\_2$ must be a Hilbert space-valued Gaussian process. What we mean by the covariance function associated is meant to be the covariance function of these stochastic processes/random operators.

Distinction between $C\_{B\_1}$ and $c\_{B\_1}$

Initially, it is important to note that $B\_1$ maps functions landing in $\\mathbb{R}^{J\_1}$ to functions landing in $\\mathbb{R}^{J\_2}$. Consequently, as detailed in Section 2.2, lines 89-91, applying its covariance function to two input functions, $\\mathbf{f}$ and $\\mathbf{g}$, yields $\\mathbf{C}\_{B\_1}[\\mathbf{f},\\mathbf{g}]$. This result is a function that takes two domain elements and outputs a $J\_2 \\times J\_2$ matrix. This matrix structure accounts for potential cross-covariance between the output dimensions of $B\_1[\\mathbf{f}]$ and $B\_1[\\mathbf{g}]$. Therefore, when we state that $\\mathbf{C}\_{B\_1}[\\mathbf{f},\\mathbf{g}] = \\mathrm{c}\_{B_1}[\\mathbf{f},\\mathbf{g}]\\boldsymbol{I}\_{J\_2}$, we are indicating that this matrix is proportional to the identity matrix. This condition is equivalent to requiring that the output dimensions of $B\_1[\\cdot]$ are independently and identically distributed (i.i.d.) for all input functions.

Assumptions on the activation function.

We would like to clarify that some assumptions on the activation function are necessary for any neural operator to be mathematically well-defined, including finite-width one, thus are included when we assume $Z\_J$ is a NO. For example, Kovachki et al. (2023) assume measurable, linearly bounded activation functions, a condition slightly weaker than Lipschitz continuity, which includes activations like ReLU, ELU, tanh, and sigmoid. Nonetheless, our statement in line 150 (Section 3.1) is overly restrictive and could be misleading; it should say, "if $\\sigma$ is Lipschitz [...] then this is a well-defined operator," rather than "only if." We will amend the main text and add to the supplementary material a proof that, for a toroidal domain, any measurable linearly bounded activation induces a well-defined $L^2$ operator.

As for the functional form of the activation function, closed-form solutions are only known for certain activation functions but these do not invalidate the result theorem. However, it is not limited to just ReLU and sigmoid. We have cited these solutions because they are the most popular activations, but the infinite-width neural network literature has continued to explore more closed-form solutions and numerical approximations to Eq (16). For example, Han et al. (2022) [arxiv:2209.04121] provides ways to compute exact and approximate solutions for a wide class of activations, and these have been implemented in the neural-tangents Python package.

What is the purpose of the regression experiment in section 5.2?

It is true that Section 5.2 does not relate directly to our results in Section 3, however, it shows that, regardless of the model's width, the model’s band-limit needs to at least cover the data’s band-limit for accurate predictions and that, even as the width increases, the predictive performance of finite-width models is not directly related to that of infinite-width ones.

Thank you for attentively reviewing and appreciating our work. We hope our answers below appropriately address your concerns. If this is the case, please consider raising your score. Otherwise, we will be more than happy to discuss our contributions in further detail.

How is the infinite width neural operator implemented in practice?

As discussed in Theorem 3.1, an infinite-width neural operator is equivalent to a Gaussian process with a specific kernel $c\_\\infty$. In Section 4, we show how to compute this kernel in two cases: Eq. (32) and Eq. (36). Given this, we can implement the infinite-width NO just like any kernel method. For instance, given some training and test data, we can compute the Gram matrix $\\boldsymbol{C}$ of the kernel by evaluating it on every pair of training points. We can sample from the corresponding GP by multiplying independent Gaussian samples with the Cholesky decomposition of $\\boldsymbol{C}$, and the inverse of $\\boldsymbol{C}$ can be used to compute the posterior for prediction.

Additionally, our supplemental material includes an implementation of the infinite-width model in PyTorch.

Assumptions behind the theoretical results.

We apologize for the possible confusion regarding the lift and projection layers. Kovachki et al. (2023) separate the layers of their architecture into three sections: lift, NO layer, and projection components. However, as we discuss in Section 2.1, both the lift and projection operators can be represented by just the NO layer by simply setting $\\boldsymbol{K}$ to zero. As such, in Theorem 3.1, when we defined $\\boldsymbol{K}\_\\ell$ and $\\boldsymbol{W}\_\\ell$ to map functions from $J\_{\\ell-1}$ channels to $J\_{\\ell}$ channels, this includes the possibility of either of them being zero. Our theorem requires that all $J\_\\ell$ values approach infinity, encompassing all layers, both $\\boldsymbol{K}$ and $\\boldsymbol{W}$, thus the hidden layers of the lift, NO layer, and projection components all go to infinity.

As for the number of Fourier modes, our theorem is applicable to any number of Fourier modes, regardless of whether they are sufficient to represent the data. As stated in Theorem 3.1, the only requirements are that the weights and kernel for each component operator are independent and identically distributed (i.i.d.) and that their covariance shrinks with increasing width. In the FNO model, Section 4.1 demonstrates that the number of Fourier modes (2B+1) acts as a configurable hyperparameter in kernel computation. Additionally, our theorem is general enough to encompass different NO architectures which might not even have a limited number of Fourier modes, such as the Matérn kernel of Section 4.2 which inherently possesses an infinite number of Fourier modes, though in practice, it is truncated by the input data's bandlimit.

Dear reviewer, thank you for your feedback. Below, we address your concerns. We hope our answer elevates your appraisal of our work. Please let us know if you have any other questions or need more clarification. We would be happy to engage further.

Do the authors think it worth discussing Deep Gaussian Processes [1]? They seem a related method that applies in the Bayesian inference setting, and therefore potentially a better comparison than finite-width FNOs.

To the best of our knowledge, DGPs have not yet been applied to operator learning and it is not clear how it could be done. In terms of our theoretical results, previous work on neural networks indicates doing either infinite width limits or infinite depth and width limits recovers Gaussian process models, but not DGPs. Our current paper generalizes these results to the operator learning setting, thus, we did not find a link with DGP models.

Do the authors see any routes to applying these methods in practice, given they perform less well than finite width FNOs in this simple regression setting? Do they see any impact on finite width FNO design that could be gleaned from this paper?

The purpose of the regression experiment is not to evaluate the performance of FNOs; it is rather to stress-test the FNOs when the true width and bandlimit are violated in the model specification and show that their performance is not correlated, as our paper is not focused on advocating for infinite-width FNOs as a modelling tool. A priori, this does not mean that infinite-width FNOs are not applicable in more realistic scenarios.

However, based on previous results on infinite-width neural networks, we see promising avenues of application of our work. For example, Adlam et al. (2023) [arxiv:2303.05420] has successfully scaled infinite-width NNs on the millions of data points setting and achieved competitive results against SotA methods and Meronen et al. (2020) [NeurIPS 2020, arxiv:2010.09494] derived activations for finite-width NNs based on infinite-width models, achieving better out-of-distribution detection and uncertainty quantification.

Can the authors explain why lower J seems to return better results for finite FNOs? This seems unexpected to me.

We speculate that a possible reason behind the performance decrease of FNO as the width (J) increases is that our training set has n = 100 functions, giving a total of (2*5+1)*100 = 1,100 data points, which is much less than the number of parameters.

Thank you for your valuable input and pointing out places where our text can be improved and other typos. We hope our answers sufficiently address your concerns. If we have left blank spots or you would like additional clarifications, we will be happy to discuss further.

Connections with existing FNO literature, further theoretical results.

Thank you for highlighting these recent empirical and theoretical results on FNOs. Indeed, we believe that the conclusion of Neal (1995), that, compared to finite-width Bayesian neural networks, their infinite-width counterparts are "disappointing", has become the common sense in the community, as it has been shown in both the Bayesian setting and in the empirical risk minimization setting (e.g. neural tangent kernels) that such infinite-width networks do not learn features from data. This analysis has been expanded by the muP / mean-field perspective, where it has been shown that there exists an optimal way to initialize networks to induce feature learning in the so-called “edge of chaos”.

As Neal (1995) laid the foundation for the analysis of infinite-width neural networks, our paper establishes the sufficient conditions for the existence of the infinite-width limit of NOs and characterizes their covariances at initialization. While we are very interested in continuing this analysis into infinite-width NOs trained by stochastic descent, similar to neural tangent kernels, we couldn't have done so without this foundation.

Regarding connections between our paper and the works of Gerog et al. (2024) and Qin et al. (2024), we would not expect our theoretical results to shed light on questions about optimal width or band-limits of FNOs, in the same way that we do not expect NN-GPs to be very illuminating about the specific post-training behavior of finite-width NNs.

Why a J=1 ground truth FNO was investigated, whether an infinite-depth perspective explains why increasing FNO depth worsens performance, and if Eq. 26 sheds light on the limitations of the infinite-depth FNO in Figure 3.

We chose J = 1 as the ground truth to stress-test the infinite-width models, as we would expect this model to be the most different from the infinite-width neural operator, as shown in Section 5.1, compared to other values of J.

Additionally, we would like to clarify that our experiments only consider one-hidden-layer deep FNOs and infinite-width FNOs. That said, we speculate that a possible reason behind the performance decrease of FNO as the width (J) increases is that our training set has n = 100 functions, giving a total of (2*5+1)*100 = 1,100 data points, which is much less than the number of parameters. Nonetheless, we would like to recall that the main point of Figure 3 is showing that, regardless of the model's width, the model’s band-limit needs to at least cover the data’s band-limit for accurate predictions and that, even as the width increases, the predictive performance of finite-width models is not directly related to that of infinite-width ones.

Claims around experimental results

Thanks for pointing this out. We agree that we have oversold our experimental results. As suggested, we will soften the tone of lines 11 and 35 in our revised manuscript. To move our results toward a more realistic setting, we set up an additional experiment with the same set-up as in the paper but using the 1D Burgers' equation dataset with $\\nu = 0.002$ from PDEBench. The task is to predict the end state (t=2) given the initial condition (t=0). Due to memory constraints while inverting the kernel matrix, we subsample the total dataset data to n = 100 functions and a grid size of 103. Our preliminary results are as follow:

At a glance, the finite-width FNO has a different optimal width for this experiment and still there is no correlation between results from models with increasing width and the test loss from the predictive posterior mean of the infinite-width model.