A Structured Matrix Method for Nonequispaced Neural Operators

We start by thanking the reviewer for your appreciation of the merits of our paper and your welcome suggestions to improve it. Below, we address the concerns raised by the reviewer and thank the reviewer in advance for their patience in reading our detailed reply.

1. We start by addressing the reviewer's concerns about framed contribution of the article and presentation around the "structured matrix method": In this regard, we would like to point out a possible misunderstanding about our contributions. While we agree completely with the reviewer that there has been considerable body of work on non-uniform FFTs and related algorithms in the signal processing literature, among others, our aim in the paper was NOT to propose a new method for evaluating the discrete Fourier transform efficiently on non-uniform grids. Rather, our goal was to propose a computationally efficient realization of widely popular neural operator architectures, such as FNO, which are currently restricted to uniform grids, on arbitrary point clouds. To this end, we needed to replace the FFT module inside the Fourier layers of FNO with an alternative that can be computationally efficient on unstructured grids. We do so by leveraging the following: i) the observation, backed by both empirical work as in the original FNO paper as well as more recently by theory (see ArXiv:2210.01074 and ArXiv:2304.13221), that FNO type models only require a fixed small (compared to the number of points) number of Fourier modes for accurate operator approximation ii) the DFT structure, being a special Vandermonde matrix, makes it relatively straightforward to obtain a fast algorithm by directly evaluating exponential sums iii) availability of highly optimized functions within PyTorch such as torch.bmm and torch.matmul for tensor operations that allows us to computationally realize the fast DFT on arbitrary point clouds. We combine all these ingredients into the SMM algorithm and demonstrate through considerable empirical evaluation that it provides a simple, accurate and fast realization of neural operators on arbitrary point clouds. This is the main contribution of this paper and it has been acknowledged as such by other reviewers. We have modified the main text to more explicitly frame our contributions and to mention that we do not propose a new non-uniform FFT algorithm but rather employ a suitable one within the framework of operator learning. This should also be contrasted with existing algorithms that extend neural operators to arbitrary grids. These approaches rely on either interpolation to regular grids or learnable geometric transformations between point clouds and regular grids, making them difficult to train and generalize and/or inaccurate. We also demonstrate through our experiments that our approach leads to a significant improvement over the state of the art methods in this regard (see Table 1).

2. With respect to the reviewers' comments on the experiments, comparisons, and existing NUFFT libraries: we thank the reviewer for pointing this out. Following your excellent suggestion, we have now compared our approach with existing NUFFT libraries. Of course, to apply these neural operators to nonequispaced data, any type-1 NUFFT could be a valid approach, as it is the FFT which limits the neural operator to the equispaced regime. However, we require not only a GPU implementation but also support within the NUFFT for backpropagation in order to train the neural operators. Therefore, the PyCuFiNufft of the referenced library cannot be used for this purpose. Given this, we opt for the TorchKbNufft library, which offers both the Kaiser-Bessel and Toeplitz NUFFT algorithms with support for deep learning architectures. We have repeated the Burgers' experiment on the nonequispaced data with these NUFFT algorithms and present the results in the Table below (see also Table 1 of the main text).

The results show that our approach clearly and significantly outperforms these NUFFT algorithms in both speed and accuracy. The KB NUFFT uses an interpolation step before applying the FFT, and the Toeplitz NUFFT corrects imperfections in the frequency domain by using a scaling matrix. Because these algorithms will be called several thousands of times within the training process, it is more efficient to do the interpolation beforehand, as opposed to in the loop. Likewise, quicker algorithms like the Toeplitz NUFFT will introduce errors greater than those of interpolation. These observations might very well have been the reason why NUFFT approaches have not been considered in the context of operator learning before and the preference has been towards interpolation approaches. Nevertheless, we thank the reviewer again for asking us to compare with NUFFT as we believe that it further strengthens the rationale and contributions of our paper. We have also explicitly mentioned this point in the main text of the revised version now.

3. With regards to the reviewers' comments on terminology and transform representations, we start by acknowledging that some terms were indeed used loosely and imprecisely in our original manuscript. In particular, the adjoint and the inverse are indeed not equivalent in the nonuniform case. This was an abuse of terminology on our part. For uniform discretizations, the inverse DFT is used for the Fourier neural operator in the discrete case, and we chose terminology to reflect this practice. To clarify the underlying process, we compute the product of an integral kernel matrix and the data in the frequency space, where we then use the conjugate transpose of the Vandermonde-structured matrix to transform our data back from the frequency domain to the spatial domain. The Vandermonde-structured matrices have an extremely ill-conditioned inverse, which is why explicit computation of the inverse of Vandermonde-structured matrices is rarely done. Thus, to be more clear to the reader, we have now changed the term "inverse transformation" to the "backward transformation" to denote the mapping back from frequency to physical space.

To maintain clarity, we also have corrected a minor typo in the section "Extension to N-dimensional Point Clouds." We intended to describe the particular example of the two-dimensional point cloud of $n$ points. In this case, the data should be represented as $X\in \mathbb{R}^{n}$ or $\mathbb{C}^{n}$, while the positions are stored as $P \left[ \boldsymbol{p}_0, \boldsymbol{p}_1 \right]\in \mathbb{R}^{n \times 2}$, the first and second columns being associated with the positions of point along the first and second spatial axes, respectively. This is not to say that the samples are reduced to simply 2 points, but rather that we have $n$ sample points whose positions are defined by their locations along the two spatial dimensions.

Furthermore, we agree that there are nicer ways to write such transforms, such as through summation notation, which is often done in literature. However, as the implementation of this transform within the context of machine learning through a PyTorch library is of special interest, we choose to describe the matrices themselves, as opposed to the transform. For problems where points vary between training samples, it is also paramount that these matrices may be efficiently constructed, mainly by avoiding for loops. Therefore, we hope to provide a general illustration of the structure of the matrix itself, as well as the implementation available in the provided code so that others planning to implement FNO-based methods on nonequispaced point-clouds will be able to tailor such an approach to their specific application.

Since I am not sure if a notification gets sent out upon the update of a review, please note that I have updated my review after reading the authors response.

At the outset, we would like to sincerely thank the reviewer for their constructive feedback as well as for raising their rating for our paper. The reviewer is rightfully concerned about the framing of our contributions and we take this opportunity to further clarify our rationale for the choices that we had made in the original version of our paper. As we have stated in the paper and as the reviewer (and other reviewers) have acknowledged that our primary contribution is indeed the realization of FNO (and related neural operators) on arbitrary point distributions, a topic of great importance in the learning of partial differential equations. Clearly, our paper should not be construed as a contribution in the general area of structured transforms and we do apologize if our framing of the title, abstract and related work section led to a perception of contributions to structured transforms. We would like to rectify this misperception as much as possible in the CRV, if accepted. To this end, following the discussion with all the reviewers, we propose to change the title of the CRV, if accepted, to Fourier Neural Operators on arbitrary point clouds or Neural Operators on arbitrary grids. This will directly frame as contribution in terms of operator learning, rather than to structured transforms. Consequently, we will modify the abstract as well as the introduction, related work and discussion to very clearly state our exact contribution as we have framed it in our rebuttal to your evaluation. We will pay special attention to your comments on the related work section in this regard.

It is also pertinent here to point out that our choice of the title was motivated by the need to differentiate our paper from papers such as "Nonequispaced Fourier PDE Solvers" and "Geometry-Aware FNOs", which claim to realize FNO on arbitrary point clouds but do so via either interpolation or learnable transformations between irregular and regular grids. This is NOT what we do, so we chose a title that reflects this difference but we believe that the proposed modified title will reflect our true contributions in a much sharper manner.

Finally, the reviewer rightly points out the use of FFT and Spherical harmonics as red herrings. In the original versions of FNO, the idea to use FFT was to reduce computational complexity as the authors thought that a large number of modes are necessary. This thinking has been persistent in the community ever since although many people have observed empirically that in practice, very low number of modes actually performs on par or better than large number of Fourier modes and the expressivity in FNO comes from large lifting layers, residual connections, and nonlinear activation functions. Theoretical understanding of this phenomena is much more recent, see ArXiv:2210.01074 and ArXiv:2304.13221. But it is precisely these observations that have motivated our approach which based on direct sum evaluation and leads to a very simple method for extending FNO to arbitrary point clouds, which is both more accurate as well as more computationally efficient than state of the art baselines. One could argue that this observation of ours was quite straightforward but to the best of our knowledge, we were first to do so and also to implement it at scale. This constitutes the core of our contributions in this paper and we aim to state it even more clearly and explicitly in a CRV, if accepted.

We thank the reviewer again for your very constructive feedback that has helped us to better frame our contributions and significantly improve the quality of the paper. We hope to have addressed your concerns satisfactorily through this reply and if so, we kindly request you to upgrade your evaluation accordingly.

We start by thanking the reviewer for your appreciation of the merits of our paper and your welcome suggestions to improve it. Below, we address the concerns raised by the reviewer and thank the reviewer in advance for their patience in reading our detailed reply.

1. At the outset, we would like to thank the reviewer for pointing out that our method is essentially the realization of FNO/SFNO on arbitrary grids, something that the literature has promised but never delivered so far. We fully agree with your assessment in this regard. It is also true that we could have mentioned this point much more explicitly in the main text of our paper. One reason why we did not do so was to emphasize the even broader scope of our method. It is not limited just to FNO/SFNO, but in principle, any neural operator that uses spectral transformations (Fourier/Spherical Harmonics/Wavelet, etc.) can be efficiently realized on arbitrary point clouds using our structured matrix method. Already in the paper, we have demonstrated this with FNO, SFNO, UFNO, and FFNO (See Table 1), and the extension to other neural operators of this type is straightforward. However, we agree with the reviewer that FNO (and SFNO) are the most important neural operators of this class in practice. Hence, and following the excellent suggestion of the reviewer, we have explicitly (and repeatedly) pointed out in the revised version that our method realizes FNO/SFNO on arbitrary point clouds. We hope that this explicit message resonates much more with the end users of neural operators in scientific machine learning.

2. Regarding the reviewer's questions and comments about how the proposed method compares to FNO/SFNO on regular grids and putting our method in the right perspective in this regard: we would start by pointing out we had already provided such a comparison for the one-dimensional Burgers' Eqn (the very first two rows of Table 1) where we had presented the results for FNO on a uniform (equispaced) grid with the underlying Fourier transforms being computed by FFT and compared with the Fourier transforms being computed with SMM. From this table, we see that both methods are completely comparable in terms of training time as well as accuracy. Following your suggestion, we have performed an additional experiment for the 2-D shear layer with data sampled on a (dense) uniform grid for FNO, with FFT and SMM realizing the Fourier transform, respectively. We find that FNO with FFT has a per-epoch training time of $251$ secs with an $L^1$-test error of $6.76\%$. On the other hand, FNO with SMM has a per-epoch training time of $211$ sec and an $L^1$-test error of $6.6\%$. This clearly reinforces the conclusions of the 1-D experiment by demonstrating that, on a regular grid, SMM is completely comparable to FFT as they are simply different computational realizations of the Fourier transform inside FNO. We have now modified the main text of the paper to clearly point out this fact.

Finally for the Shallow water equations on the sphere, we followed the reviewer's suggestion and trained SFNO (as well as FNO) on the underlying regular grid on a sphere. However, we evaluated both methods on randomly distributed points on the Sphere (outputs of SFNO and FNO can be evaluated on any point in the domain), and the resulting test errors are now shown in SM Table 5, as well as directly below. Comparing with the results obtained with SFNO-SMM (and FNO-SMM), we see that the errors are very similar whereas the SMM algorithm is faster to train in this case. With these additional results, we hope to have addressed your concerns about putting the results in the correct perspective by comparing with FNO/SFNO on regular grids.

3. Regarding the reviewer's comments and questions about how sensitive is the proposed method to the choice of collocation points: the reviewer is correct in pointing out that, in analogy with classical numerical methods, one can expect larger errors on unstructured grids, in general. This is indeed borne out by the experimental evidence already presented in Table 1 for the 1-D Burgers' Eqn example. In the table, we present results on a regular equispaced grid and on an irregular expanding-contracting point distribution (See SM Figure 4) to observe that the test error with FNO-SMM increases by almost a factor of $4$ when transiting from a regular to an unstructured grid. We have added a further comment on this issue in the main text and thank the reviewer again for pointing it out to us.

We start by thanking the reviewer for your appreciation of the merits of our paper and your welcome suggestions to improve it. Below, we address the concerns raised by the reviewer and thank the reviewer in advance for their patience in reading our reply.

1. The reviewer's point about introducing more information on FNO is very well taken. Following your excellent suggestion, we have now added an entire section SM A.1, where we explicitly define the FNO. We have also referenced this section in the main text. However, given the space constraints, we could not incorporate a full section on FNO in the main text itself.

2. In response to the reviewer's point about it being hard to visualize the proposed method is bringing about in the whole neural operator learning workflow in Figure 1, we would like to say that we are primarily providing a more efficient method (SMM) that replaces the FFT and inverse FFT of the Fourier layer of the FNO and allows FNO to be realized on arbitrary point distributions. Other than that, the workflow is unchanged, and the FNO diagram would be identical. We have now explicitly mentioned this fact in the main text of the revised version of our paper.

We hope that we have addressed the reviewer's concern to your satisfaction. If there are any further questions we could answer to upgrade the reviewer's confidence in their assessment of our paper, we would gladly do so.