Operator Learning with Neural Fields: Tackling PDEs on General Geometries

Summary

This paper creatively employs Implicit Neural Representations (INR) for operator learning on irregular domains. This novel method benefits from INR's capability to adaptively handle irregular grid distributions or irregular geometric areas, making learning the mapping from input to output in the INR's latent space quite natural. The primary experiments of the paper are on initial value problems, where it shows a significant improvement. However, a notable challenge is that the introduction of INR necessitates bi-level optimization, which potentially introduces instability and difficulty into the training process.

Strengths

1.The proposed method is highly innovative. 2.The quality of writing is excellent, and the paper is easy to read. 3. As an empirical paper, the experimental results surpass the baselines.

Weaknesses

1.The paper lacks research on related work. Besides Geo-FNO and FFNO, transformer-based methods can also naturally handle irregular geometric areas. The authors should at least mention these works. Here are some references for consideration [1,2,3]. 2. Operator learning tasks are not computationally demanding, so the paper should explore the method's stability and report error bars (at least for some cases). 3. More explanation or supplementation is needed concerning the stability of the bi-level optimization, the difficulty of training, and why second-order meta-learning is adopted. Overall, this is a good paper. If the authors can clarify the issues mentioned, I am inclined to accept it. Reference 1. Transformer for Partial Differential Equations' Operator Learning (https://arxiv.org/abs/2205.13671) 2. GNOT: A General Neural Operator Transformer for Operator Learning (https://arxiv.org/abs/2302.14376)

Questions

None

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

None

Summary

This paper proposes a method for solving PDEs with neural networks in continuous space (and optionally time as well). The authors leverage the success of coordinate-based neural networks (or implicit neural representations – INRs) and formulate the problem as an operator learning one, akin to the Neural Operator family, bypassing the need for discretisation. In particular, the method is based on a typical encode-process-decode strategy: 1. *Encode*: the input condition of the PDE (i.e. the initial value in an initial value problem – IVP) is projected to the input latent vector (via gradient-based optimisation), which is then transformed (via an MLP) to a latent modulation of the input INR (SIREN), such that the INR is fitted to the input condition (auto-decoding). 2. *Process*: the input latent vector is transformed (via an MLP when projecting to a single output or via a neural ODE when solving over continuous time) to the output one. 3. *Decode*: the output latent vector is transformed (via an MLP) to a latent modulation of the output INR, which encodes the output signal and can then be queried in arbitrary spatial positions. To promote training stability, the authors employ a two-step process: First the parameters of the input/output INRs are fitted via a meta-learning strategy, similar to Dupont et al., ICML’22, using a shared INR across all input conditions and learning both the parameters of the INR and the latent vectors simultaneously and efficiently. Then, the parameters of the encode-process-decode pipeline are learnt by doing optimisation in the latent space. The method is evaluated on several PDE tasks, showing competitive results against other continuous neural solvers.

Strengths

- **Clarity and presentation quality**. The paper is well-presented, has a good flow and guides the reader smoothly throughout the important ideas. Also, it is mathematically clean and concise. Finally, it is mostly well-contextualised with the relevant literature. - **Simplicity and Scalability (inference)**. Although the method is less theoretically motivated compared to Neural Operators, it builds on top of existing tools in the INR literature, making it easy to implement. In addition, forecasting happens in a compressed version of the physical process (a latent vector), which improves scalability. - **Novelty and Significance**. The paper proposes a new framework for continuous simulations, departing from the Neural Operator family, advancing the field further, in terms of efficiency and performance. The claims are empirically evaluated on a broad spectrum of PDE tasks (from IVPs and dynamics modelling to geometric design and optimisation) showcasing competitive generalisation performance across input conditions, resolutions and time horizons (for trajectory unrolls).

Weaknesses

- **Connections with Dupont et al., ICML’22**. First off, it is important to acknowledge that the method, from a technical point of view, is largely based on the “functa” concept of Dupont et al., ICML’22. The authors do cite this work, but not as prominently as one would expect. Although the aforementioned paper focuses more abstractly on function representation learning, and less on PDE solvers, the current work is a natural follow-up, so I believe that the connections between the two should be stated more clearly. - **Efficiency and ease of training**. The training process might be slow since it happens in two steps. Although the second step, might be relatively fast since it happens in the latent space, the first one requires fitting INRs to all possible inputs and outputs in the training set. The meta-learning strategy might potentially alleviate this problem, but the authors have not discussed this matter. Could the authors provide more details on this step? How long does it take? Does it require extensive hyperparameter tuning? How sensitive are the results to different hyperparameters? - **Potential risk of overfitting (w.r.t. input/output signals not coordinates)**: Experimentally, it seems from the paper that the method generalises well across different input conditions, including the particularly challenging problem of trajectory unrolls. Let me explain why I am a bit surprised by this: - I am puzzled about the inductive biases of the encoding process (Eq. (1)). In particular, it is unclear what are the hidden assumptions regarding the input/output signals. Think of a CNN encoder for example – in that case, we would silently assume signal stationarity and shift invariance (due to the application of the same filter across different regions), locality, etc. On the other hand, the corresponding assumptions of the auto-decoder are unclear, even in an intuitive manner. I would appreciate it if the authors could comment on this, and it might be also useful for future readers. - Given the above, I am wondering how well the encoding process, and more importantly the entire pipeline, generalises across different input signals. In particular, in case the encoding process overfits, then small errors in the predictions of the latent codes (especially when unrolling trajectories where errors accumulate), might result in regions of the latent space “unknown” to the decoder, and therefore to unpredictable outcomes. To clarify this, it might be useful to discuss the distributions of the input conditions used. Are they challenging enough? How important is the Z-score normalisation (Appendix B.1.2.)? - My concern about this also comes from the fact that only a latent-space objective is minimised, which in principle does not rule out the possibility of small errors in the latent space to translate to large errors in the signal space (e.g. assume that the INR decoder has a large Lipschitz constant). - **Empirical evaluation and comparison against relevant methods**: I think there is room for improvement in the experimental section (extra discussion and comparisons) in order to make the claims more convincing. In particular, - It would be helpful to also include discrete neural solver baselines (e.g. based on CNNs, as in Stachenfeld, ICLR’22, or mesh GNNs, as in Pfaff et al., ICLR’21). - In addition, it would be nice to see a more thorough discussion regarding the comparison between continuous neural solvers (MPPDE, Neural Operators etc.). It is unclear to me why and when the current method performs better. For example, intuitively I don’t understand why MPPDE is dependent/overfits on the training grid. Have the authors considered comparing against the Graph/Multipole Graph Neural Operator methods (NeurIPS’20)? Does CORAL have any stronger inductive biases compared to Neural Operators? Could the authors elaborate on the above? - Regarding the dynamics modelling experiments, I have a few questions: (1) Have the authors compared against autoregressive one-step predictors (e.g. as in the MPPDE paper, or in Sanchez-Gonzalez et al., ICML’20)? Obviously, the NeuralODE formulation allows for evaluation across arbitrary timestamps, but I am wondering if there is any disadvantage in this approach. (2) Have the authors considered comparing to discrete autoregressive solvers, such as the LE-PDE paper (NeurIPS’22)? (3) In certain papers (e.g. MPPDE and Sanchez-Gonzalez et al., ICML’20) the authors use a noise trick to stabilise trajectory unrolls and mitigate error accumulation. I am curious if the authors used it here as well, and if not why?

Questions

The major questions that I would like to be clarified and discussed are included in the "Weaknesses" section. Below, I enlist a few extra minor comments: **Minor comments**: - Could the authors explain their inference-time findings? Why is CORAL faster than DINO and MPPDE and slower than DeepONet and FNO? Why does CORAL scale batter (almost constant time regardless of the resolution)? Why didn’t the authors include results for FNO across all resolution setups? - Perhaps it would have been beneficial to also test on more challenging setups such as the turbulent regime of Navier-Stokes (potentially to identify limitations and failure cases). - A recent paper that improves upon prior work and might be useful to also include as a baseline is the “Clifford neural layers for PDE modelling”, Brandstetter et al., ICLR’23. - L72: “infinite-dimensional functions“: This should be rephrased - maybe mappings between functions (infinite-dimensional vectors)? - L76: Regarding DeepONet, if I am not mistaken it is required to have the same observation grid only for training and not for testing. - L265: RK4 – maybe explain that you are referring to fourth-order Runge-Kutta.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The authors have included a section discussing the limitations of their work and were upfront about them. A few things remain to be clarified (e.g. time and tuning required for training, generalisation across input conditions, and more challenging experimental PDE setups). No foreseeable negative societal impact.

Summary

This work tried to introduce an implicit neural field-based framework for solving PDEs for different applications on general geometries. The method represents the input and output function spaces as an implicit neural representation. The method consists of a two-stage pipeline: first, the authors train two auto-decoders over the input and output function domains. Then, using paired data, they learn a mapping between the spatial latent code between domains. The paper shows results in initial value problem, dynamics, and predicting PDE solution from different geometry. I think the paper is tackling a very important problem - generalizing operator and neural field method to solving PDE with different geometry (in the boundary/shape). And I also like the approach to this problem as they leveraged neural fields techniques to achieve discretization-free effect and operator techniques to allow relatively fast influence time. However, the presentation of the paper confuses me in several places (see weakness section) and that holds me back from full acceptance in the first read.

Strengths

- The method is discretization free. It doesn't require the user to discretize the domain, which can have effect on the simulation quality. This method only need to run optimization with auto decoder, which is capable of taking partial observation and predict a continuous function that satisfies these observations. - This method has shown advantages compared to many operator methods in that it doesn't require operations such as Fourier transformation which requires regular samples in the geometry. This advantages the method to potentially work with different domain geometry, as well as different observation density.

Weaknesses

1. Can this method be used for Inverse design? The experiment named "geometry design" is a bit misleading as it hints at the designing the geometry (the domain). But to the best of my reading, the geometric design experiment is set-up to predict PDE solution from different geometric, rather than designing the domain geometry. It would be great if the authors can add a note on this in the supplementary. 2. It's not clear to me how does the method actually generalize to different geometry - I might miss something here. For example, how does the method encode where the boundary of the geometry is? It occurs to me that all functions are defined in the ambient space of physics domain. This makes it unclear to me how does the effect of domain boundary on the solution gets encoded. For example, if we have a bunch of samples (x_i, u_0(x_i)) that indicates the initial condition of a IVP problem, but we set two different boundary condition - one circle these samples with a sphere and the other with a rectangular box. The current method seems to make the same prediction as the latent code is inferred from the set of samples samples (x_i, u_0(x_i)), but I can't believe that's the case in the real world - i.e. the boundary has no effect on the solution. 3. Justification for the requirement of simulated data. The system still requires paired data for training similar to the neural operator, which is more supervision required by the PINN or neural fields methods (which only requires the PDE definition). I wonder what's the way to justify this additional data? 4. Risk of overclaiming. I think "general geometries" has become very vaguely define throughout the paper. I'm under the impression that the method is able to solve equation with different boundary condition and different geometry of the boundary. But it's not clear to me how the current method is able to handle different boundary conditions, also it's not clear to me how does the method take into account of the different geometry of the simulation domain. For example, if we have the same initialization function, but different shape, would the latent code we obtained from test-time optimization of the auto-decoder be the same or different? If in test time, we want to change the boundary condition of this problem, how do I specify this? In general, I think the proper claim of the method is that it allows the neural operator method to work with irregular samples, but every trained model might still be restricted to solving the same set of problem (maybe the same boundary geometry) but with different initialization.

Questions

See weakness section.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Section 5 discussed limitation

Summary

In the paper, the author proposed to use neural field in operator learning. In the CORAL model, it first encodes the input into some codes, apply an MLP on the codes, and then decode into the output function. Since the neural field can be continuous evaluated, the CORAL model can be applied to general geometries.

Strengths

The paper uses neural presentation in operator learning which is a very natural choice and worthy investigation. The proposed method show promising results on time extrapolation and geometry design problem.

Weaknesses

The overall framework is quite simple. The encode-decode structure has been studied in PCA basis [1] and Fourier basis [2], which are also applicable for general geometries. Compared to these well-studied conventional basis, neural representation usually use a higher number of parameters, which in the end, make the encoding unstable. In the CORAL framework, the latent code is obtained using gradient descent. However it's possible to get very different latent code given the same input. Would such non-uniqueness cause any issue? [1] Bhattacharya, Kaushik, et al. "Model reduction and neural networks for parametric PDEs." The SMAI journal of computational mathematics 7 (2021): 121-157. [2] Fanaskov, Vladimir, and Ivan Oseledets. "Spectral neural operators." arXiv preprint arXiv:2205.10573 (2022).

Questions

What is the dimension/size of the latent code? One of the major advantage of this CORAL is time extrapolation. Do the authors have some insight why neural representation does better on this task? In the experiment, the authors compared CORAL with another neural representation based method DINO. How does CORAL compared to DINO in the model design? Does it also use encoder and decoder? It will be great to address DINO also in the related work section. In neural representation, many recent work use context grid and voxels to enhance the representation. Is it possible to compensate CORAL with grid-based representation?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The paper did not address its limitations.

Thanks for the response!

Thanks authors for the response! The replies address some of my confusion. In general, I think the paper proposes a system that combines the advantages of neural field representation with advantages of neural operators. This can be of interests to NeurIPS audiences. I will keep my acceptance rating.

Discussion

Thank you once more for your helpful comments and insightful recommendations. As we approach the discussion deadline, we kindly want to remind you of our keen interest in hearing your response and discussing the points raised in the rebuttal. We would like to emphasize that we have included additional experiments in the rebuttal to address the questions raised in the various reviews. Thanks again, The Authors

Response to Authors

Dear Authors, Thank you for your reply. Most of my concerns have been addressed. I will keep my score unchanged and vote for acceptance, as per my initial review (well-presented work, simple to implement, works well in practice, a new framework for continuous simulations). However, I would like to point out that the underlying reasons that lead to the success of this method, as well as its limitations, are still unclear to me. For example, the authors do mention that in the turbulent regime of NS the performance deteriorates substantially, but they have not discussed what are the input condition distributions that are currently tested, so as to understand the complexity of the tasks to be solved. Moreover, since there is no theoretical evidence and the inductive biases of the INR auto-decoding are hard to grasp, I think that there is a considerable gap in our understanding. I would encourage the authors to try to provide more intuition in an updated version (e.g. including the intuitive explanation given in the “Potential risk of overfitting” paragraph of their rebuttal) and investigate this further in the future. Some minor concerns that were not addressed and would be good to discuss in an updated version: - Looking at Table 2 it is apparent that the INR training creates a large computational burden in the overall pipeline, with the processing in the latent space being only a small fraction of the overall training time. Could you compare the overall training times with those of the baselines? - It is claimed that training is relatively robust to the hyperparameters. I would suggest adding some quantitative results in an updated version and discussing how the hyperparameters are chosen (is it by looking at the training error?). - The authors have not discussed the following questions: - “To clarify this, it might be useful to discuss the distributions of the input conditions used. Are they challenging enough? How important is the Z-score normalisation (Appendix B.1.2.)?” - “My concern about this also comes from the fact that only a latent-space objective is minimised, which in principle does not rule out the possibility of small errors in the latent space to translate to large errors in the signal space (e.g. assume that the INR decoder has a large Lipschitz constant)” - Could the authors clarify what “exponential sampling” is?