Scalable Transformer for PDE Surrogate Modeling

Summary

The paper presents a high-dimensional function decomposition technique for mitigating the complexity encountered by Transformers when dealing with high-dimensional data. The authors have developed a clear logical argument and a comprehensive review of the relevant literature. It conducts extensive experiments on several benchmark 3D problems governed by NS equations. It uses latent marching which shows better empirical performance compared with autoregressive models. However, the paper's novelty appears to be somewhat limited, particularly as similar ideas were presented at ICLR the previous year, including Factorized Fourier Neural Operator (FFNO) and Tensorized FNO. Overall, the paper lies on the borderline. It presents a clear motivation and technical soundness, and has shown some empirical advantages. However, the limitations highlighted above, particularly the lack of comprehensive experimental comparisons and some confusing notations, hinder its overall quality. I lean towards accepting this paper if the authors could strengthen the experimental results and address the raised concerns.

Strengths

1. The motivation for the research is clear, with the Introduction section being easy to read and understand. 2. From a technical perspective, the paper appears sound, with a well-articulated methodology that would be easily replicable, especially if the authors decided to open-source their code and datasets. 3. Empirically, the proposed method has shown some advantages.

Weaknesses

1. The experimental section is somewhat weak. The authors considered 3D datasets only within regular geometric areas and did not compare their approach with many of the available baselines such as U-FNO, FFNO, HT-Net, Tensorized FNO, and Wavelet NO. The paper's baselines only included standard FNO and ResNet, making it challenging for readers to ascertain whether the Transformer-based approach indeed excels in high-dimensional data processing. I suggest the authors incorporate additional representative baselines (at least one or two) during the rebuttal period. 2. The use of notations in the paper is a little confusing. Does 'd' represent the dimension of the hidden layer? How is it selected? It seems like it would be beneficial to annotate some of the intermediate variables' dimensions or shapes, as the current description in the paper is not very clear. 3. In equation (7), the 'v' kernel appears to still increase exponentially with the dimension, which begs the question, how is this overcome the curse of dimensionality? 4. While the proposed method seems to involve some sort of tensor decomposition on the data, the paper does not investigate the relationship between these processes in detail. It would be beneficial to explore this theoretically, as it could shed light on the model's representational capacity. After all, certain tensor decompositions have theoretical guarantees for approximating a high-dimensional tensor. 5. I suggest you to include the following closely related work in the reference if you have not mentioned them in the paper. FFNO, HT-Net and Tensorized FNO are proposed in the last year to deal with high dimensional data. WNO, U-FNO are also popular architectures for operator learning. And you might need to consider citing a more relevant survey that list many related works on neural operators. Ref: 1. Factorized Fourier Neural Operators (https://arxiv.org/abs/2111.13802) 2. Multi-Grid Tensorized Fourier Neural Operator for High Resolution PDEs (https://openreview.net/pdf?id=po-oqRst4Xm) 3. HT-Net: Hierarchical Transformer based Operator Learning Model for Multiscale PDEs(https://openreview.net/forum?id=UY5zS0OsK2e) 4. Multiwavelet-based Operator Learning for Differential Equations (https://arxiv.org/abs/2109.13459) 5. U-FNO -- An enhanced Fourier neural operator-based deep-learning model for multiphase flow (https://arxiv.org/abs/2109.03697) 6. Physics-Informed Machine Learning: A Survey on Problems, Methods and Applications (https://arxiv.org/abs/2211.08064)

Questions

None.

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

None.

Summary

This paper uses separable axis attention to reduce the complexity from being exponential in spatial dimension to linear.

Strengths

NA

Weaknesses

The standard Transformer has two components: attention and MLP. This paper only addresses the attention. But the MLP complexity O(Nd^2) would still make the overall complexity exponential in spatial dimension. This seems to defeat the purpose of the paper. By design, the attention without softmax has rank d, where d is the hidden size and is typically much smaller than the sequence length N. This is shown in Fig 8a. For A^(1) and A^(2) in Fig 8b and 8c, the sequence length along each separate axis is much smaller. This gives the impression that matrices A^(1) and A^(2) are less rank deficient. But this does not mean that the proposed attention is better than vanilla attention.

Questions

Table-4 only reports the encoding or attention layer runtime. What is the overall runtime including MLP layers?

Rating

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

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

3 good

Presentation

2 fair

Contribution

3 good

Limitations

NA

Summary

This paper proposes a new method to scale Transformer-based models for PDE surrogate modeling to higher dimensions. The new method combines several "tricks" from previous works, and most importantly, uses factorization to further reduce the computational cost of training and evaluating Transformer-based operator learners, while not sacrificing the accuracy too much, and even achieves superior performance in certain benchmark problems. Even though the idea of using factorization (tensor decomposition) for end-to-end problems that feature a spectrum decay and/or sparse structures is not new, I check the Transformer-based PDE modeling papers and this is the first time someone implemented it. I think this is a worthy contribution to the literature, however, there are apparent weaknesses in the paper as well (detailed below). I lean toward acceptance if the author can address or answer some of the concerns below.

Strengths

- The authors proposed a simple way to reduce the computational cost of linearization of attention further. The method is actually quite simple by using a different projection and iteratively evaluating the attention "integral", which is essentially tensor decomposition when the un-integrated kernel matrices are moved to the innermost integral written together. As the reason behind tensor decomposition makes more mathematical sense than the one featured in Axial Transformer. - Dramatic savings versus even the linear attention. - I actually appreciated that the authors spilled some techniques that are usually deemed "tricks" explicitly in Section 3.3, and used them across all baselines, instead of hiding them in the source code as some works on Transformers. - The authors enhanced many existing PDE benchmarks (in the anonymized GitHub repo), and created several new ones that test the scalability of the models.

Weaknesses

- The biggest weakness is perhaps some lack of theoretical foundation, but I guess this is fine for a methodology paper. - The saving of computation in the presentation of formula (7) may not be that obvious to the community of Transformer research, who may not be that familiar with the integral representation. Especially considering the presentation of (7) uses $n$-D, but compares that in (2), which is for a 1-D problem. - The presentation in laying out the factorization could use some improvement for people more familiar with linear algebra than differential equations. For example, on line 179, - If only the author could add a more informative diagram comparing the difference between the factorized approach vs Axial Transformer. ### Misc small typos - line 144: there should be space between RoPE and its reference. - line 152: this is just a suggestion, the tilde notation could be clearer if it is referred back to equation (3), saying something like "$\tilde{\mathbf{z}}$ denotes a matrix $\mathbf{z}$ whose row vectors $\mathbf{z}_j$ are RoPE encoded as in (3)". - line 154: "learnable projection" -> "learnable projections". - line 209: in (9), $Z$ should be set to $\operatorname{Att}(U)$.

Questions

- In the comparison of SVD for attention matrices, what is the value of $k$? What specifically are $A^{(1)}$ and $A^{(2)}$? From different model problems? - In Section 4.2, the FNO models have significantly higher parameter counts. Is there a specific reason for this?

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

2 fair

Contribution

3 good

Limitations

The author properly acknowledged the curse of dimensionality. Moreover, the dimension-bound factorization limits the domain type to tensor-product type (i.e., rectangular).

Summary

The proposed work presents a scalable transformer architecture for modeling partial differential equations (PDEs). The input function is linearly projected to multiple functions with one-dimensional domains using a learnable integral projection operator. The attention mechanism is then applied to these projected functions. This factorized attention mechanism requires less time and memory. The proposed FactFormer shows comparable or better performance while requiring fewer parameters and less memory.

Strengths

1. The proposed method is effective and scalable, making it suitable for high-resolution PDE modeling. 2. The paper is well-written, and the method is described in detail.

Weaknesses

1. The experiments in this work use fluids with very high viscosity, resulting in simulations of flows with low Reynolds numbers. As the model uses linear projection operations and performs softmax-free attention in the projected functions, it is crucial to demonstrate the method's effectiveness in modeling complex systems, such as flows with high Reynolds numbers. For instance, the set of experiments performed in FNO to model the Navier-Stokes equation with varying viscosities.

Questions

1. is there any non-linearity or activation function used in the architecture? 2. How well does the model perform in the zero-shot super-resolution task compared to the FNO baseline?

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

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

This paper proposes Factorized Transformers (FactFormer) for surrogate modeling of partial differential equations (PDEs). FactFormer first projects the input function to multiple one-dimensional projected functions, which are then used for multi-dimensional factorized attention. The factorized attention greatly reduces the model complexity, while still achieving competitive performance on three turbulence problems.

Strengths

1. Improving the efficiency of transformer is important to the development of transformers in PDE modeling. 2. The proposed FactFormer achieves superior model efficiency with competitive performance.

Weaknesses

1. The design of learnable projection and factorized kernel integral may limit the model capacity, especially for problems that do not have such underlying low rank structure. For example, the projection block basically gets $n$ projections of the input function on $n$ dimensions, which could drop the crucial information about the input function. 2. In principle, full attention should have better model capacity than factorized attention. However, the experimental results in this paper show that full attention is much worse. The authors hypothesize that it is due to the instability of rollout. An error trend plot of full attention and factorized attention is need to verify this hypothesis. 3. The authors hypothesize that CNN variants are better than FNO on turbulence problems because CNN filters capture high-frequency patterns but FNO truncates the high-frequency modes. However, the nonlinear activation function and the local linear transform branch in FNO allows FNO to model high-frequency patterns. [1] has shown it can efficiently approximate operators in incompressible Navier-Stokes equations. Further justification and ablation study is needed to support this hypothesis. 4. The inference speed of FactFormer is 1.5-1.8x slower than FNO while the number of parameters is much smaller. Why is that? Can you also report the FLOPs of the FactFormer and compare it with other baseline models? 5. Overall, the accuracy of FactFormer is not consistently better than Dil-ResNet. Minors: 1. Line 179: $L_N$ should be $L_n$. 2. Line 302: "time cost for ..." -> "time cost of ..." 3. Line 315: "time cost by ..." -> "time cost of ..." [1]: Kovachki, Nikola, Samuel Lanthaler, and Siddhartha Mishra. "On universal approximation and error bounds for Fourier neural operators." The Journal of Machine Learning Research 22.1 (2021): 13237-13312.

Questions

1. How is the error trend plot computed exactly? I thought it is accumulative error but the error is not monotonically increasing. 2. Can you also compare the model efficiency with tensorized FNO, which leverages the tensor factorization to improve efficiency of the original FNO? 3. Section 4.3 claims that Factorized attention has higher ranks so it is more efficient. Can you elaborate more on how you jump to that conclusion?

Rating

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

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

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Listed in weaknesses and questions sections.

Summary

The authors propose FactFormer, a neural PDE solver that is based upon the transformer, but scales much better with the dimensionality of the PDE problem. The authors compare FactFormer to several other neural PDE solvers on three different problems and find that Factformer achieves similar performance to state-of-the-art PDE solvers, but while offering lower computational costs. This property may also allow FactFormer to scale to high-dimensional problems than existing PDE solvers.

Strengths

1) The paper is generally well-written 2) The authors use three challenging PDE problems, and they compare their approach to two state-of-the-art existing models, providing reasonably good empirical evidence for the performance of their approach. 3) The authors do provide some analysis comparing FactFormer to other Attention mechanisms, which is helpful for proving their advantage.

Weaknesses

1) (Major) Upon inspection, it doesn't appear that these baseline models (FNO and Dil-ResNet) have been previously applied by other authors to any of the three benchmark problems (with the particular settings used in this work), so that the authors here had to implement and apply the baseline methods themselves. This introduces a major potential source of bias if the authors, because it is unclear whether the authors made equal effort to optimize the performance (e.g., by adjusting model size, batch sizes, learning rates, etc) of the baseline methods and their proposed approach. It would be better if the authors could test FactFormer on one or two PDEs where the baseline methods (FNO and Dil-ResNet) have previously been optimized and tested by other authors, helping to ensure a more fair comparison. For example, why don't the authors use the data/benchmarks from [84], for example? 2) (Major) The main claim of this paper is that attention is good, but we cannot use it because of the computation time. Therefore we make it more efficient, resulting in FactFormer, but the FactFormer model doesn't perform any better than existing models (e.g,. Dil-ResNet). Therefore, if there was an advantage to using Transformers for PDE solvers, apparently that advantage disappears when using the factorization in the Factformer. So then the main contribution here is presumably the computational time of Factformer? Factformer does seem to be significantly more computationally efficient than Dil-ResNet for Inference - which is great - but what about training time? If it is not also superior during training then this is a major limitation. 3) (Minor) The main stated contribution of this paper is to extend transformers to solve higher-dimensional problems, rather than outperform other PDE solvers. Indeed, the Dil-ResNet often outperforms the FactFormer, and therefore there seems to be little, if any, performance advantage against recent PDE solvers. The authors show some theoretical big-"Oh" analysis of runtime, but these are asymptotic bounds. In Table 4 the authors report some computation times, but it is just for one problem, and it is unclear what problem they are testing on (perhaps I missed it?). Since this is a main claimed contribution of this work, it seems quite important to me to report the empirical computation time (e.g., per epoch or iteration) for PDE problems with some varying mesh sizes and dimensionalities, so that we can see how the runtimes compare for some common settings of these parameters.

Questions

1) Per Weakness #2 above, What is the computational cost of training FactFormer compared to Dil-ResNet and FNO? 2) Is there a reason the authors did not use existing benchmark PDE problems to test FactFormer (per my Weakness #1 comment above)? It is important to me that there is a good reason, or alternatively, that the authors can either (i) provide evidence that they fairly optimized all competing models (this is very hard to do), or (ii) add another benchmark that was used in a prior study, and compare Factformer's performance to that reported for the baselines on that prior study. If these questions are (convincingly) addressed I am willing to significantly increase my score, potentially to "Accept".

Rating

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

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

na

Thank you for your detailed response, and quick followup questions

I thank the authors for their detailed response to my comments. I find most of the argumentation and new results convincing. I have increased my rating by one point I will consider increasing it by another point or two, but I have one more follow up question regarding the fair comparison of FactFormer with prior competing approaches. In Table 1, the authors added additional experimental results for the Navier-Stokes and Darcy-Flow problems, which I appreciate, and which I believe invariably strengthens the results of the paper. Per my previous commentary however, these results would be even more convincing (much more in my opinion) if the results of the competing methods (e.g., FNO, Dil-Resnet) on this problem were reported from prior work, where the authors of the competing methods had optimized these models for the problems under study. With this in mind I have three (hopefully quick ) follow up questions: 1) Is that still the case that you (the authors) implemented the competing models yourself on these new problems (Navier-Stokes and Darcy)? 2) If yes, then why? e.g., is the data for these problems from prior studies not publicly available so you could run FactFormer on the same data? 3) Even if data from prior studies is not available, how do your results with the competing models compare with those obtained from prior studies? Are they similar? If the authors could provide a short description and/or point me to a particular study, and the relevant Table/Figure numbers, that would be helpful.

Response to the Rebuttal

I thank the authors for their response. The proposed method also achieves comparable performance in modeling fluid flow with high Re. I am increasing the score.