Latent Neural Operator for Solving Forward and Inverse PDE Problems

Summary

The present paper introduces a novel architecture, coined Latent Neural Operator (LNO), made of a sequence of self-attention layers intertwined between two Physics-Cross-Attension (PhCA) encoder and decoders. Experiments conducted on multiple PDE datasets evaluate the efficacy of the proposed approach across various metrics and highlight its potential advantages.

Strengths

The incorporation of self-attention layers in neural operators for solving PDEs has recently received significant attention. This paper contributes to this evolving field by introducing a novel architecture that integrates self-attention mechanisms within a stack of transformer blocks applied in the latent space. This design choice appears to enhance computational efficiency. A notable strength of the paper is its comprehensive evaluation, with the inclusion of an ablation study.

Weaknesses

- *Originality*: The incorporation of self-attention layers in neural operators is not new and is well-documented. To this regard, multiple references are missing. The primary novelty of this work appears to lie in the integration of a stack of transformer blocks with self-attention mechanisms within the latent space. - *Quality*: The paper's main contribution is the introduction of a novel architecture evaluated through numerical experiments. However, the absence of theoretical guarantees represents a notable weakness. It is also very concerning that the authors have copied and pasted Table 1 from "Transolver: A Fast Transformer Solver for PDEs on General Geometries," and merged it with their results. This practice raises questions about the originality and integrity of the numerical experiments presented. See also my concerns in the "Limitations" review section. - *Clarity*: The paper is difficult to read and lacks a coherent structure, especially in the numerical experiments section. Note that, since the contribution is mainly about the proposed architecture, a detailed description of the selection of the hyperparameters would have been appreciated. Overall, it is not very clear how this work differs from related works and what it borrows. Understanding the content requires multiple back-and-forth readings. The authors should focus on organizing their content more logically and clearly differentiating their contributions from existing research to enhance the readability and impact of their paper. In the same line, the paper contains numerous grammatical errors, which hinder its readability. I highly recommend that the authors thoroughly review and correct these errors to improve the clarity and overall quality of the paper. - *Significance*: While the paper suggests there may be some performance gains, it lacks clear demonstration of the significance of its findings in advancing the state of the art. A more rigorous comparison, such as cross-validation of the architecture and ensuring fairness in number of learning parameters with existing methods, would mitigate its somewhat limited contribution. Please also refer to my comments under the "Quality" and "Limitations" sections.

Questions

- The learning tasks in Section 4.2 are not entirely clear to me. Could you formalize each task and specify the number of samples available for each? - I have some concerns about the scaling experiments of Section A: are the width and depth selected based on the L2 test errors? How do you perform neural architecture search for state-of-the-art methods? - Why does the number of parameters change from dataset to dataset in Table 4?

Rating

4

Confidence

4

Soundness

2

Presentation

1

Contribution

2

Limitations

The paper does not adequately provide the statistical significance of the experiments. I recommend that the authors include a thorough statistical analysis to support the validity and reliability of their experimental results.

Summary

This paper suggests a neural operator architecture for learning solutions of forward and inverse PDE problems using a sequence of transformer layers. The novelty of the architecture is that after a carefully designed initial transformation step, named “Physics-Cross-Attention”, the inputs to subsequent transformer blocks have dimensions completely independent of the number or dimension of the samples of the input function. This design is meant to allow for efficient processing of data at very fine resolutions. The paper evaluates the efficacy of their method on standard PDE forward problems discretized on regular and irregular grids; they report superior performance compared to 11 other neural operator methods. The paper also considers two inverse problems dealing with Burgers' equation. One problem is an interpolation of randomly-sampled data, and another is extrapolation to recover initial conditions given sparsely-sampled data.

Strengths

Overall, the paper is relatively strong in terms of originality, quality, and significance. I believe this is a technically strong paper that requires editing to become a strong submission. I will list individual strengths below: - The proposed “Physics-Cross-Attention” module is a simple innovation of the popular cross-attention mechanism. I think it is likely others will build off this idea due to its easy implementation. - The experimental results on the forward PDE problems are impressive. When compared to 11 other methods, this method outperforms all others on 4 out of the 6 problems considered. - The solution method for the inverse problem is interesting and novel. The solution method breaks down the problem into one of first interpolating randomly-located samples to a regular grid and then extrapolating the reconstruction on a regular grid to earlier and later time points. If this is a novel solution method, I think it is valid to claim this as another contribution, and if this is inspired by other works, readers would be interested to know where this solution method comes from.

Weaknesses

The clarity of the paper could be greatly improved. I will list individual weaknesses below: - The paper claims to solve PDEs in a novel way by transforming from a "geometric space" to a "latent space" and solving the PDE in the "latent space". These terms are never defined, which makes it difficult to distinguish this paper's innovation from past work. The “Physics-Cross Attention” module transforms an input of shape $(N_{in}, d+1)$, where $N_{in}$ is the number of sample points and $d$ is the dimension of the spatial variable, to an array of shape $(M, D)$, where both $M$ and $D$ are hyperparameters completely divorced from $N_{in}$ or $d$. However, other neural operators have an encoding step that transforms inputs of shape $(N_{in}, d+1)$ to shape $(N_{in}, p)$. In this case, one could say the samples are projected to some latent space in $\mathbb{R}^p$. Clearly defining the meaning of “latent space” considered in this paper would better define the novelty of this work. - The paper claims that the hyperparameters $M$ and $D$ are independent of the number of discretization points in the input $N_{in}$ and the number of discretization points in the output $N_{out}$. This is not supported by the experiments, which do not vary $N_{in}$ or $N_{out}$ at test time. Without an evaluation of the model’s robustness to varying $N_{in}$ or $N_{out}$ at test time, it could be the case that $M$ and $D$ were tuned to have good performance on datasets with the specific settings of $N_{in}$ and $N_{out}$. - The description of the model architecture is confusing due to notational choices. I have added specific questions about the architecture below. - The experiments in sections 4.1, 4.2, and 4.3 are missing some important details. I have added specific questions about experimental details below. - The paper claims this method speeds up training time by 1.8x when compared to the state-of-the-art Transolver architecture. This is not supported by the relevant table (Table 4), which only reports time per epoch and not the total number of epochs needed to train each model. - There are grammatical mistakes throughout the text. These mistakes did not obscure the meaning of the text but were quite distracting. These mistakes could be addressed by software such as Grammarly or an AI chatbot; using these tools for grammar checking is allowed by NeurIPS 2024.

Questions

Questions about the introduction: - The introduction alludes to the connection between the attention mechanism and kernel integral operators. Can the Latent Neural Operator (LNO) architecture be interpreted as a new way to parameterize kernel integral operators? If so, is the new parameterization motivated by the physics or numerics of the problems considered? I believe this type of discussion would better motivate the architecture choices made, and the paper would have more impact if the architecture choices were justified in this way. - The introduction discusses the improved computational complexity of this method over other methods. Presumably, there are specific problems that the LNO method could solve but are intractable for other methods. Could the authors give examples of such problems? This could help to better motivate the paper. Questions about the architecture and notation: - What are the input and output dimensions of the branch and trunk projectors? This confusion is caused by conflicting notation. In section 3.1, $x$ is a spatial location, but in section 3.2, $x$ is the output of the trunk projector. Similarly, $D$ is the spatial domain in section 3.1 but is a shape hyperparameter in Figure 1. - In Figure 2, what is the white box labeled V? - What is the intuition guiding the softmax being applied across different dimensions in Figure 2? - The figures make it seem like the "Trunk Projector" is separate from the "Attention Projector", but these are both MLPs, and they are always used in sequence. Why make the distinction between these two networks? Could one think about them as one single MLP in the "Trunk Projector" module? Questions about the experiments: - In Table 1, what is the meaning of the two groups of comparison models separated by a horizontal line? - In Table 1, what is the randomness in the 5 independent trials? Different train/test splits? Different initializations of the model weights? - In Section 4.1, for each benchmark task, what is different between the train and test set? - In Section 4.2, what exactly are the inputs and outputs of the model for the completion experiment? Similarly, what exactly are the inputs and outputs for the propagation experiment? Are different initial conditions used? - In Figure 6, it appears like the solution of Burgers’ equation gets smoother as it progresses through time; Burgers' equation is known to form shocks as waves propagate through time. If this plot is correct, can the authors comment on why there are no shocks forming in this example?

Rating

7

Confidence

4

Soundness

2

Presentation

1

Contribution

3

Limitations

The authors identify one limitation, that the model does not scale well to deeper architectures (i.e. more transformer blocks), in the NeurIPS Paper Checklist, but I can not find this limitation discussed in the text or appendix.

Summary

The paper introduces Latent Neural Operator for solving forward and inverse problems in latent space. It uses a Physics-Cross-Attention to encode the input to the latent space, uses self-attention to evolve the state in the latent space, and uses another Physics-Cross-Attention to decode back to the real-world space. Experiments in forward and inverse problems demonstrates the effectiveness of the method.

Strengths

The paper addresses the important problem of learning PDEs and improving its efficiency. The paper is mostly written clearly.

Weaknesses

1. In terms of novelty, the paper is somewhat limited. The usage of cross-attention to embed the input into latent tokens is much like the one in LSM [1], with the main difference being that the paper uses self-attention in latent evolution. In addition, the main modules of the paper are all based on attention, which is a standard neural architecture. 2. The paper lacks comparisons to important past works in latent space, including forward problems of [2,3], and inverse problem of [4]. 3. There are some questions that needs to be resolved in the Questions section. References: [1] Wu, Haixu, et al. "Solving high-dimensional pdes with latent spectral models." ICML 2023. [2] Iakovlev, Valerii, Markus Heinonen, and Harri Lähdesmäki. "Learning Space-Time Continuous Latent Neural PDEs from Partially Observed States." Advances in Neural Information Processing Systems 36 (2023). [3] Wu, Tailin, Takashi Maruyama, and Jure Leskovec. "Learning to accelerate partial differential equations via latent global evolution." Advances in Neural Information Processing Systems 35 (2022): 2240-2253. [4] Zhao, Qingqing, David B. Lindell, and Gordon Wetzstein. "Learning to solve pde-constrained inverse problems with graph networks." ICML 2022

Questions

1. Why does in Table 1, FNO and GeoFNO gets the exact same result in Darcy and NS2d? 2. Why does F-FNO achieves worse performance compared to FNO in NS2d? In the original paper, it achieves much better performance than FNO in NS and etc.

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

The paper can benefit from a more in-depth discussion on its limitations.

Summary

The paper presents the Latent Neural Operator (LNO), a novel approach for solving forward and inverse partial differential equations (PDEs) by operating in a latent space. The LNO introduces a Physics-Cross-Attention (PhCA) module for transforming data from the geometric space to a learnable latent space and demonstrates improved prediction accuracy and computational efficiency, reducing GPU memory usage by 50%, accelerating training by 1.8 times and achieving state-of-the-art accuracy on several benchmarks.

Strengths

1. **Clean model design**: The model design looks very clean, achieving the separation of positions and corresponding values so that they can be input into different model modules. 2. **Excellent result**: The proposed model surpasses nearly all models currently in this field, achieving SOTA.

Weaknesses

1. **Lack of theoretical proof**: Although the results are fantastic, can a set of theories support the current structure design? 2. **Lack of real-world data examination:** Many 3D datasets exist in this area, such as smoke and turbulence[1]. Adding these benchmarks would strengthen the model's efficacy claim. Industrial car simulation data can also be used to test the proposed model's performance in this paper[2]. [1] Li, Z., Shu, D., & Barati Farimani, A. (2024). Scalable transformer for pde surrogate modeling. *Advances in Neural Information Processing Systems*, *36*. [2] Wu, H., Luo, H., Wang, H., Wang, J., & Long, M. (2024). Transolver: A fast transformer solver for pdes on general geometries. *arXiv preprint arXiv:2402.02366*.

Questions

1. It is a good idea to separate the position and the corresponding value. Still, for the physics-cross-attention proposed in the paper, only the position is used to generate the attention map. Will this operation cause a lack of information? Their correlations may not be the same for the same points in different initial fields. 2. Can all operations be done in the latent space? The decoder decodes the latent information only when it needs to generate an output and is not involved in the autoregressive process. 3. Does the model in Table 4 share the same parameters as the one in Table 1? It needs to be a fair comparison. 4. I think the term decouple is not very rigorous because, in a branch projector, the positions and corresponding values are used simultaneously, so they are not completely separated.

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

Lack of real-world benchmarks and some theoretical proofs.