Scalable Transformer for PDE Surrogate Modeling
Authors: Zijie Li
Dule Shu
Amir Barati Farimani
Venue: NeurIPS 2023 poster (2023)
Keywords: Efficient attention
Neural PDE solver
TL;DR: We propose a multidimensional factorized attention mechanism that alleviates the curse of dimensionality when applying to PDE problems.

Abstract: Transformer has shown state-of-the-art performance on various applications and has recently emerged as a promising tool for surrogate modeling of partial differential equations (PDEs). Despite the introduction of linear-complexity attention, applying Transformer to problems with a large number of grid points can be numerically unstable and computationally expensive. In this work, we propose Factorized Transformer (FactFormer), which is based on an axial factorized kernel integral. Concretely, we introduce a learnable projection operator that decomposes the input function into multiple sub-functions with one-dimensional domain. These sub-functions are then evaluated and used to compute the instance-based kernel with an axial factorized scheme. We showcase that the proposed model is able to simulate 2D Kolmogorov flow on a $256\times 256$ grid and 3D smoke buoyancy on a $64\times64\times64$ grid with good accuracy and efficiency. The proposed factorized scheme can serve as a computationally efficient low-rank surrogate for the full attention scheme when dealing with multi-dimensional problems.

Scalable Transformer for PDE Surrogate Modeling
Zijie Li, Dule Shu, Amir Barati Farimani
Carnegie Mellon University Mechanical Engineering Department
{zijieli, Hulleseaaddre. cmu. edu & barati@cmu edu
Abstract
Transformer has shown state-of of-the-art performance various applications and has recently emerged as a promising tool for surrogate modeling of partial differential equations (PDEs). Despite the introduction of tinear-comllexty attention, applying Transformer to problems with a large number of grid points can be numerically unstable and computationally expensive. In this work, we propose Factorized Transformer (FactFormer), which is based on an axial factorized kernel integral. Concretely, we introduce learnable projection operator that decomposes the input function into multiple sub-functions with one-dimensional domain. These lub-functions are then evaluated and used to compute the instance-based kernel with an axial factorized scheme. We showcase that the proposed model is able to simulate 2D Kolmogorov flow 256 x 256 grid and 3D smoke buoyancy on a 64 x 64 x 64 grid with good accuracy and efficiency. The proposed factorized scheme can serve computationally efficient low-rank surrogate for the full attention scheme when dealing with multi- dinteeiimnsii problems.
1 Introduction
Various physics processes are modeled by partial differential equations (PDEs), from the interaction between atoms in molecular systems to large-scale cosmological phenomena. Solving PDEs advances the understanding of complex physical phenomena, enabling people to make accurate predictions, and make informed decisions across wide range of scientific and engincering disciplines. Numerical solvers provide a practical way to simulate and predict PDEs since many PDEs are often difficult to solve analytically. Most numerical solvers divide the continuous domain into discretized grid and reduce the continuous differential equations to algebraic equations via methods like finite dffferenccellemee//vllmm methods or spectral method. Despite the theoretical guarantees behind them, their practical realization of specific problems can pose challenges that require careful expertise overcome, such sufficient understanding of the underlying physics, resolves the necessary spatio-temporal scales. The interest in developing user- seer-frienlly and efficient PDE solvers, along with the success of deep learning models many other areas [13, 44, 51, 103], has facilitated the emergence of neural-network- herlllnnttork-bbseee -based PDE solvers, where the neural network can be used to parameterize the solution function of the target equation [96], or approximate the solution operator [69, 76]. Compared to many numerical solvers, neural PDE solvers appear to be more tolerant with coarse discretization [104], and can be applied without explicit meshing [96]. In addition, knowing the underlying equations are not strictly necessary for neural PDE solvers, which gives them the potential to simplify and accelerate the process of physics simulation based on PDEs.
fine-tailored mesh that
Among various neural network designs, attention-baser models (Transformer) [115] have become state-of-the-art for wide array of applications [13, 18, 27, 51], which gives rise to a recent surge of interest in applying the Transformer to PDE modeling [16, 29, 35, 41, 43, 54, 67, 73, 82, 86]. By viewing the input sequence a function sampled on a discretization grid, attention can be interpreted as a learnable kernel integral [17, 35, 54, 60] or a learnable Galerkin projection [16], and the saaeeeeeeto-sequenee Transformer [115] can be modified correspondingly be better suited for
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Pred.
Ref.
Figure 1: Model's prediction (pred.) and reference ground truth (ref.). Left: 2D Kolmogorov flow on 256 x 256 grid; Right: 3D smoke buoyancy on 64 64 x 64 grid (zOy cross-section is shown).
PDE modeling [16, 35, 43, 54, 67, 74, 86]. In these works, attention is typically applied to every grid point in the domain to exploit both the local and non-local structure of the system, and therefore a linear-complexity variant of attention is usually necessary. As the number of grid points grows exponentially with respect the number of dimensions, this results in very large attention matrix that computes the interaction between every pair of the grid points (despite this attention matrix is evaluated explicitly in linear attention). Consequently, cascading deep stack of attention layers introduces instability and relatively high computational cost on nigh-resolution grid. To alleviate these issues and improve the scalability of Transformer in PDE modeling, we propose a modified attention mechanism. Our model is inspired by the kernel integral viewpoint of softmax-fin -rre attention, with a factorized integration scheme motivated by the inherent low-rank structure of dot-product kernel matrix. More specifically, we propose a muutti-dimensioall factorized kernel integral with each kernel function in the integral having only dingle-dimession axial kernels, we propose learnable integral operator that is able to project the input function with nigh-dimensiooal domain into a set of sub-functions with single-dimensional domain. The computation of each axial kernel is quadratic with respect to the number of grid points along that the corresponding axis but does not grow with the number of dimensions, which alleviates the curse of dimensionality in standard attention. With the modified attention mechanism, scale up to ciati--dimnsionall problems with a large number of grid points and achieve competitive performance compared to state-of- the-art art models. Moreover, we show that our factorized attention mechanism can reduce the computational cost and improve stability compared to softmax-fit linear attention.
domains. To calculate these
proposed model can
2 Related works
Neural PDE solver Based on the emphases of model design, neural PDE solvers can generally be divided into the following groups. The first group of work focuses on using neural networks with mesh-specific architecture design (such as convolutional layers for uniform mesh, graph layer for irregular mesh) to learn the spatial and/or temporal correlation of the PDE data [10, 38, 48, 64, 65, 75, 88, 91, 92, 100, 104, 109, 114, 116]. With input-target data collected, the training process can be conducted without the knowledge of underlying PDEs. This can be appealing when the physics of the system is unknown partially known, such as large-scale climate modeling [61, 83, 90, 97]. The second group of work, namely the Physics-Informen
Neural Networks (PINNs)[15, 40, 42, 52, 77, 87, 96, 106, 127], treat neural networks a parametrization of the underlying solution function. PINNs incorporate the knowledge of the governing equations into the construction of loss function, which includes the residual of the PDE, the consistency with given boundary condition and initial condition. Unlike the previous group of works, PINNs do not necessarily need input- target data and can be trained solely based equation loss. The third group of works, often referred to as the neural operator, focuses on learning a mapping between the function spaces[5. 8, 9, 16, 36, 43, 50, 54, 60, 68, 70, 71, 76, 78, 86]. Neural operator has the generalization capability within a family of PDE and can potentially be adapted to different discretization without retraining. DeepONet [76] proposes a practical realization of the universal operator approximation theorem [21]. Concurrent work graph neural operator [69] proposes learnable kernel integral to approximate the solution operator of parametric PDEs and the follow-up work Fourier Neural Operator (FNO) [68] achieves excelling accuracy and efficiency on certain types of problems. Broadly speaking, the
'Code for this project is available https: //t::///gitt /ttgtttttt.t///////////ete...het www.th....oo..aaaeeeeeeeeeeee.ee cwmc

---

operator learning can be conducted upon different types of function bases, such as the Fourier bases [34, 58, 59, 68, 95, 110, 121], wavelet bases [36]. learned bases in an attention layer [16, 67], or based on approximation of the Green's 's function [7, 108). The training of neural operators can also be combined with the principle of PINNs to yield more physically consistent prediction [72, 117]. Our model is closely related to the neural operator, as the major building blocks in our proposed model are a learnable projection operator and a learnable kernel integral operator.
In addition to direct surrogate modeling, neural networks can also be combined with numerical solvers to improve their accuracy and efficiency. For example, using a trained neural network to correct the error of the solver on the fly [3, 28, 56, 89, 113], or doing offline high-fidelity reconstruction [24, 30, 49, 66, 102].
Transformer for Physics Simulation The Transformer model [115] have gained outstanding popularity in natural language modeling [13, 25], imagery data processing [27] and beyond [51]. In the field of physies simulation, Transformer has drawn increasing research interest as surrogate model for simulation, with its modeling capability demonstrated both as neural PDE solver [16, 32, 35, 41, 43, 48, 54, 67, 82, 86] and as a pure data-driven model in the absence of known governing PDE [14, 20, 31, 83]. The dot-product attention can be considered as an approximation of an integral transform with non-symmetric learnable kernel function [16, 17, 35, 54, 60, 122], which relates Transformer to other popular operator learning models such the ENO [68]. We will expand the discussion of Transformer under the kernel viewpoint in Section 3.
Efficient Transformer Following the introduction of Transformer [115], various works have investigated ways of reducing the computational cost of standard scaled-dot product attention. The first line of work seeks to remove the softmax and make use of matrix associativity to derive linear complexity attention [23, 53, 101], which has also been explored for PDE modeling[16, 43, 67]. The second line of work tries to approximate the dot product between query and key matrix by exploiting the low-rank structure of it [4, 22, 45, 55, 120, 123, 126]. Our work is related to the first group of works with a softmax-free design, but still calculates the dot product between query and key first. Among the second line of work, Axial Transformer is closely related to our work, as both works have explored conducting attention in an axial fashion. However, the derivation of attention matrix is different in the two works (see Section 3.2 for detailed comparison). More generally, the exploitation of the multi-dimensional tensor structure in our proposed model can be related to tensor factorization methods [57, 85] and their applications in various deep learning models [58, 63, 80, 84, 124].
3 Method
3.1
Attention mechanism
Standard attention Given three sets of vectors, namely the queries {q keys (ki), and values {v
No, (assuming N& = No). attention mechanism [2, 33, 79, 115] dynamically computes a weighted average of the values: z, - شi h(q, ky) k, where q,, k,, v; E I h(-) is the weight function that determines the contribution of a specific value to the final output. An example of h(-) is the scaled-dot product with softmax [115]: h(q,,k, = exp (q.kg/ /T) Σ. exp (q:k. /T), and T is usually chosen as r = Vd. The queries/keys/vauues are usually obtained from inputs via learnable projection. In self-attention, all of them are computed from the same source as follow:
q: u, W,W, k; = u, W., V; = W.W
(1)
where u, E Ridd= is the input vector and {W;, 9: WE, W.) E Rd are learnbale projection matrices. In cross-attention, queries are derived from one input while keys and values are derived from another.
Attention as learnable integral Under the hood of PDE modeling, the input sequence to the attention layer can be viewed as the sampling of input function on the discretization grid [16, 54, 60, 67]. Kovachki et al. [60] propose that the scaled-dot product attention [115] can be viewed special case of a Neural Operator [60], where the attention amounts to the Monte Carlo approximation of the learnable kernel integral. Cao [16] further proposes two interpretations of softmax- free attention. The first is to view attention as the Fredholm integral equation of the second kind with a learnable asymmetric dot-product kernel, and the second is to view as Peterov-Galerkee
projection with learnable basis function. The softmax-free attention proposed by Cao [16] is later extended in OFormer [67], where Rotary Positional Encoding (RoPE) [105] is introduced to modulate the dot product and can be viewed as another special case of the kernel integral in Neural Operator style.
3

---

In this work, we continue on adopting the learnable kernel integral viewpoint of attention and view each channel of the hidden feature map as the sampling of a specific function on the discretization grid. Given query/key/value matrix {Q. K, V) E TR xd, their row vectors: q:/ k. vi. correspond to the sampling of a set of functions (-), (-), (-)) on grid point z;. where { ] discretizes the underlying domain. As more conerete example, the 1-th column (channel) of q represents the sampling of function (-(-) on a grid point, i.e. (q.)) (zz). Furthermore, softmax-free attention is
equivalent to the numerical quadrature of a kernel integral:
(z.)² = Σ wslq k.)(v.)' = I. 1 (211, (E())
(2)
where z; is the output vector, K (I) = Ei= (() is an instance- kernel and wa is the
quadrature weight. Understanding attention from the perspective of the kernel has been an active topic of research 23, 111, 122]. The theoretical approximation power of different kernel integrals has also been analyzed under the context of PDE learning [35, 54, 60].
Note that the above kernel does not explicitly depend on the spatial coordinates (Ix;, <). For this work, we opt for a modified kernel formulation proposed in OFormer [67]. which modulates the dot product kernel with relative position. Assuming the underlying spatial domain is 1-D (which is sufficient for our proposed model, see next subsection), given query and key vectors q k, and their corresponding spatial coordinates T,, T.g. RoPE [105] (g(::)) .): : Rixd x R + Ix is defined as:
g g(q. x3) = q, θ (zx), 9 (k). = k, θ (xz)
(3)
cos (Az,0t)
- -sin (Axz0t) cos (Ax,01)
where: θ (x:) = Diag ( Ray2(1:)): R =
sin (Ax,0t)
{1, ... d/2)) following Vaswani et al. [115] and Su et al. [105]. À is mesh- based weight that we heuristically set to 64 throughout most problems. The projection function O(-) : :R + Rd x Rdd can explicitly modulate the dot product with relative position: 9 (q. (kj,
and X, Bj are hyperparameters. 0/ is usually chosen as 10000-
= 0 xz) kz thanks to the following property of rotation matrix: (a, = Ra( - zy). To summarize, we will adopt attention mechanism in the following form for the proposed model (with modification discussed in the next subsection):
Z VQKVV
where Ō denotes a matrix whose row vectors are RoPE encoded as in (3), e.g., Q. g(: .z). w the quadrature weight with typical choice of 1/N for uniform quadrature rule, Z is the output matrix. The query/key/value matrix Q/K/V is derived from the input via learnable projections defined in (1). The matrix product QKT evaluates the kernel function k(.,,) on the discretization grid (x; NN
3.2 Multidimensional factorized attention
Compared to the standard scaled dot product attention that has quadratic complexity with respect to the length of the input sequence, the attention in (4) can enjoy linear complexity by making use of the associativity of matrix multiplication (calculate K TTV first). In PDE modeling, the length of the input sequence is equal to the number of points on the underlying discretization grid. Assuming the n-dimensional domain is discretized by S x S2 x... x S,, N points, the softmax-fre free attention in (4) will compute the kernel in (2) with the dot product of Q and K, which are N by d matrices with N usually much larger than d. The kernel matrix computed is by design low-rank as is the product of two tall and thin matrices. Meanwhile, attending large number of grid points to each other can be unstable and the linear attention has a complexity that quadratic to the channel dimension d, which can limit the scalability of the model in terms of its width. To improve the numerical stability and reduce the computational cost of the aforemennioned attention mechanism, we propose a simple yet efficient way to modify the kernel integral discussed in the previous section which is motivated by the low-rank structure of attention. Essentially, our model computes the kernel integral in an axial factorized manner instead of convolving over all the grid points in the domain.
For the following discussion, we will use tensor notation [57] to describe the operation. We
assume the data is represented on uniform Eulerian grid and can be treated as n-way tensor

---

Linear
Attention
Attention
Atlention
MLP
(.,+)
(2) (+++
x (+,
OAK Linear
anK Linear
QIK Linear
Input function
Output function
(5- 200)
(i (
Learnable projection
Figure 2: Schematic of the factorized kernel attention. Upper path: the input is transformed into the Value via linear transformation. Lower path: the input is first projected into multiple sub-functions with mie-dimensional domain. These sub-functions are then used to derive the Query and Key on each axis, and their dot products form the kernel function of the corresponding axis. The Value is iteratively updated by the kernel integral transform along each axis and finally sent to MLP.
U E Xx.. The product of it with matrix W E R'xs5 across the m-th mode will result tensor of shape S x ... 5m-1 X J X Sm-1 X... X Sn. whose elements are defined as:
8m
(U Xm W) Wiigiammhm 18 1iodd
Σ Km
Uniign-menot W Wjm-
Learnable projection The first major component of the proposed framework integral operators ( g(z), ,..., , g(m) that projects the input function u R2 + Rd into set of functions with one-dimensional domain {di1), (5(2) (.))) R Rd, which is defined as:
of learnable
(5(m) (ar (m)) = (l() (n)(a.
-/im) (wfa... -
Q 7(m) (u (u(E1. 5m -m ((n) Sm+1, kn)) 1.1....d.. 1155m-1d5am 1d5m+1 denn
where h((m)() Rdd Rd and (m) Rd Rd are pointwise learnable functions and I/( Tm Im+ ... Tn). with Jmm being the size of domain Ωnm discretized by (a:(") ) S In practice, we implement (1()) as a three-layer multi-layer perception (MLP) similar to the feedforward network in Transformer [115] and (mm)) as a simple linear transformation. When the underlying grid is uniform, (6) simply amounts to first transforming the input with pointwise learnable functions ((m)) (.), applying mean pooling over all but the m-th spatial dimension, and then applying another pointwise learnable function (m))
Factorized kernel integral Equipped with the above projection module, we now introduce our factorized kernel integral scheme. More specifically, we propose to use the following integral to replace the kernel integral in (2):
(n)
2 Ji
=So
11)
E1)
/
(2)) (2) (1))) E2).... / " Ea)u((i, dd
(1
den.
9g
Rn
where kernels {x(1) x(()) I x R- R are computed based on projected single-dimensions function along each axis, v(-): : IR" + Rdd is derived from the input function u via linear trans- formation (just the value in standard attention). Next, will discuss how the above kernel integral is implemented in practice. Using the learnable projection operator defined in (6), we can obtain (iO)...., ((m)) (É)m) E 1m55kd)) from the input U E R
where the im- th row of Ü is the evaluation of projected function ((() wi., (((m) (((m) (xr...) ). Then imm
we apply linear transformation on them to obtain the query/key matrix just as standard attention: ((m)) ((()) (m)) w?(m), where (W(m), w:(m)) E Rdxd are learnable matrices. The query and key are used to compute the kernel matrix A(m) Sm...
A (m) wm. ((m) (k(m))"
(8)
II practice often have additional mode for the channel, resulting in a 1)-way tensor.
5

---

where Wm is the mesh weight, G denotes the RoPE encoded matrix as discussed in (4), the i-th row and j-th column of (m() represents the kernel value x)(m) (x."m) ..f")). Despite (8) has quadratic complexity with respect to the grid size Sm, this an affordable cost for most of the problems where the axial grid size Sn is mostly between 64 to 512. Meanwhile, the value V E x251x..x8.. is
derived from the input via a linear transformation (i.e. (m + 1) mode product): V = U xn+1 W.. where W I is again a learnable matrix. The overall factorized kernel integral is numerically approximated with the following tensor-matrix product (Figure 2):
Z = At( = V x1 A X2 A x Xn A
(9)
In (9), the computation of all kernel is of complexity O + S .... + d), and the time complexity of a single tensor-matrix product V X A (m) is O(N Smd). After evaluating the tensor- matrix product, the output tensor Z will be sent to pointwise feedforward network f(+): : Rdd + Rd. To sum up, the update protocol of a single layer in our proposed factorized Transformer is defined as follows:
U +- f (IN (Att (U))) + U,
(10)
where Att(.) is the attention from (9), IN(-) is instance normalization [112] that normalizes each channel mnstance-wise
It is worth pointing out that the axial factorized kernel proposed here shares some similarities with the Axial Transformer proposed in Ho et al. [45], but has two significant differences despite the connection. Firstly, Axial Transformer reduces the computational cost by constraining the context of attention along each axis (e.g. a pixel can only attend to other pixels on the same row), which amounts to moving all but one axis to the batch dimension. In this way computing the axial kernel matrix is of ((NSmc) complexity (recall N = S x And its overall computation of attention relatively more expensive due to the presence of softmax. Secondly, the decomposition in Axial Transformer is not layer-wise. For example, in the first layer, the attention is conducted in a now-wise manner and then the second block will conduct attention in a column-wise, manner, whereas our model decomposes attention along all axes into tensor-matre product within every layer. We provide an illustrative example in Figure 27 of the Appendix.
S,)) instead of O d) as in our model.
3.3 Training techniques
In this subsection, we will discuss several techniques used for training the model (ineluding baselines) in our numerical experiments, In general, these techniques aim to alleviate the compounding error of autoregressive neural PDE solvers when applied to time-dependent PDEs.
Latent marching It is proposed in the Li et al. [67] that simple pointwise learnable function e(:.)) E Rd x Rx0 + Rd can be used to propagate dynamics in the latent space with a fixed time interval At: (z, t + At) layer. In practice, E is implemented as a pointwise MLP and is efficient to compute. Leveraging this technique, with one call to the neural solver, (by marching in the latent space for k steps), thus reducing the total number of calls by a ratio of h. This is in principle similar to the Temporal Bundling technique proposed in Message- Pastiigg Neural PDE solver (MP-PDE [12]. yet different in practical realization. In MP-PDE, the multi-timestep prediction is implemented as first predicting the difference in time {di, d2.... ... d& and then adding them to the input uo by a forward Euler scheme in the physical space: ûn = uo + dy At. In this work, we opt for the latent marching to predict multi-timesteps as the forward Euler scheme (in the physical space) is less stable for fluid problems with relatively large time step sizes.
(z, t) + E (2 (z, t), t). where is the output of the final attention
forward the state for multiple time steps
Pushforward Neural PDE solvers are observed to be unstable for time-dependent problems. A small error or perturbation that occurs at the beginning of neural PDE solvers' prediction, can easily result an unbounded rollout error. While there is hardly a universal method for guaranteeing their stability, wide array of techniques have been proposed to improve the stability of neural PDE solvers, such as adding physics constraints [72, 118, 119], rollout training [68] or adding random-wall noise [91, 100, 104]. For this work, we adopt the pushforward technique from MP-PDE, which amounts to rolling out the model for two steps during training and then letting the gradient only flows through the last step. This allows training the model on error-corrupted samples and promotes the stability of the model. From a practical perspective, this is straightforward to implement and also computationally much cheaper than standard rollout training.

---

4 Experiment
In this section, we will investigate our proposed model numerically on several challenging problems. Furthermore, compare our model against softmax-free free attention [16, 67]. The baseline models we compared against are Fourier Neural Operator (FNO) [68]. Factorized Fourier Neural Operator (F-FNO) [110] and Dilated ResNet (Dil-ResNet) [44, 125]. FNO has been shown to have good accuracy on wide range of PDE problems and is computationally very efficient owing to the Fast Fourier Transformation (FFT). F-FNO factorizes the spectral convolution in FNO into separate spectral convolution along different axes and adopt an improved residual connection formulation like Transformer [115]. Dil-ResNet is recently introduced by Stachenfeld et al. [104] to learn the coarse- grained dynamics of turbulent flow and has demonstrated state-of ofte-of-thhe-ar ff-thee-rtt art performance across several problems. We adopt the implementation of Dil-ResNet with group normalization from PDEArena [37]. On 2D steady-state problem where linear attention's is computational cost is affordable, we also include the result from Galerkin Transformer [16]. which uses CNN to project the function onto a coarse grid and applies linear attention on the coarse grid. The implementation details of the proposed model and baselines are available in Section A, B of the Appendix.
4.1 Benchmark problems
We first apply our model to three fluid-like systems, where the underlying physics patterns are sensitive to the spatiotemporal scale that discretization can resolve, and typically require fine discretization for classical numerical solvers. In these problems, the neural PDE solver is trained to predict the next frame (or multiple frames if using latent marching) given a context of previous frames. The number of context frames of Kolmogorov flow and isotropic turbulence is set to 10 following [68, 72], and for smoke buoyancy similar to [37]. We also consider a well-known steady-stan problem- Darcy flow, which has been studied in many of the previous works. Below we provide brief description of each problem we studied. More details can be found in Section E of the Appendix.
2D Kolmogorov flow The first example is 2D Kolmogorov flow governed by incompressible Navier-Stokes equation with a periodic boundary condition. The Reynolds number Re determines how turbulent the system will be. We adopt the setting of forced turbulence following Kochkov et al. [56] and generate the data by using the pseudo-spectral method to simulate fluid flow with Reynolds number Re = 1000. The objective is to predict the vorticity w of the flow field within an interval [to, fo + -7, where T is and t is a random starting point in the sequence. We use spatial grid of 256 x 256 and temporal discretization of At = 0.0625s (therefore is corresponds to 16 frames) to train and evaluate the model.
3D isotropic turbulence
The second example is 3D isotropic turbulence governed by incompress- ible Navier-Stokes equation with a periodic boundary condition. The major difference from the first example that the vortex stretching term is non-zero for three-dimensional flow. We use the 3D spectral simulator from Mortensen and Langtangen [81], which simulates the forced turbulence described in Lamorgese et al. [62]. For generating the dataset, we simulate a system of Taylor Reynolds number Rex 84 [62]. The objective is to predict the pressure p and velocity u from t = 0.5 to f = Is (10 frames). The model is trained and evaluated on a 60 x 60 x 60 spatial grid with At = 0.05s.
3D smoke buoyancy The third example is 3D buoyancy-driven flow, which depicts smoke volume rising in a closed domain. A similar system in 2D formulation has been studied in several previous works [11, 113]. The underlying governing equation is the incompressible Navier-Stokes equation coupled with an advection equation. The boundary condition for the smoke field is Dirichlet while the boundary condition for the flow field is Neumann. The advection equation describes the motion of smoke, which is transported along the flow field. We modify the solver from [37] that is implemented in [hiffoo[[6] to generate the data, with buoyancy factor set to 0.5 and viscosity
0.003. The
objective is to predict the scalar density field of smoke and velocity of flow u from t = 3 to t = 15s (16 frames). The model is trained and evaluated on 64 x 64 x 64 spatial grid with At = 0.75s. To account for non- non-periodic boundary conditions, we pad the domain for FNO variants and DilResNet following the original works. For FactFormer, we append simple CNN block after the model, which comprises 3-by-3 convolutional layers with zero padding.
2D Darcy flow In addition to the above time-depedendent systems, the fourth example is 2D steady-state problem from Li et al. [68]. Given the diffusion coefficient, the model predicts the steady-state flow field. The boundary condition is also Dirichlet so we adopt settings for all models similar to the 3D smoke problem.
7

---

4.2 Results and discussion
For all the models, we study two protocols of training. The first is Latent Marching with Pushforward (denote LM). The second is simply Autoregressive (denote as AR), where the model is rolled out for two steps during training. For LM models, each call to the model will output future steps. On 2D Kolmogorov flow/3D smoke buoyancy, k is set 4, and 2 for 3D isotropic turbulence. We interleave pushforward training with standard per-step training for LMM models. The relative T,2 norm is used to train and measure the error of each model following Li et al. [68]. The sequence-wise averaged error and the frame-wise error at the end frame are reported in Table 1, 2, 3. We also report the time cost of simulating sequence and the number of parameters for each model. The frame-wise error trends are shown in Figure 6. 7, 8, 9, 10 in the Appendix. The visualization of predicted samples are provided in the Section F of Appendix.
We observe that Dil-ResNet has slightly better per-frame fitting capability compared to the other models on 3D flow problems. As shown in the loss trend plots, it starts at a lower error compared to other models. This coincides with the observation in Stachenfeld et al. [104] where Dil-ResNet's performance strong on 3D fluid problems. On 2D flow, F-FNO has the best accuracy compared to other models. Interestingly, FactFormer can catch up with Dil-ResNet on 2D Kolmogorov flow and 3D smoke buoyancy when the time duration becomes longer. Yet for shorter-term prediction 3D isotropic turbulence, Dil-ResNet still has the best final accuracy. This suggests that the accuracy of long-term prediction can potentially benefit from exploiting the global structure that lies in the input. Nonetheless, compared to Dil-ResNet, FactFormer offers superior efficiency as indicated by the inference time (time cost of simulating a sequence). Since the training time is roughly proportional to the model forward time, on 3D problems Dil-ResNet generally takes 3-4 times longer to train. In terms of different training strategies, we find that AR models are less stable than multi-step training (LM) and computationally more expensive it requires more calls to the neural solver. Despite the average error varies case by case, LM models' error generally accumulates slower on the problems we studied, whereas AR models quickly blow up in some cases.
Lastly, while Dil-ResNet has shown good accuracy for 3D flow problems, its performance highly de- pendent on the training discretizations. As shown in the Figure 3, without changing model architecture, its evaluation errors increases significantly when the resolution increases, while Transformer -based models and FNO models' performance are roughly invariant to the resolution. This highlights major difference between CNN-based models and neural operators.
FNO2D
F-FNO2D
Dil-ResNet
FactFormer
Model
AR"
LM
AR'
LM
AR
LM
AR
LM
avg. error
0.3177
0.2978
0.1486
0.2453
0.8156
0.1655
0.8835
0.1734
final error
0.4423
0.4567
0.2811
0.3861
1.1692
0.3051
1.0963
0.3017
Inf. time (s) # params (M)
0.73
0.81
0.86
1.01
4.69
1.78
3.14
1.38
85.1
3.7
2.4
3.5
Table I: Evaluation results of 2D Kolmogorov flow. A batch size of i0 used for inference. LM models predict 4 steps with each call to the model. Total prediction length is 16 steps. AR": Since for 2D problem FNO variants can afford to rollout more steps during training, AR FNO rollout for 12 steps, AR F-FNO rollout for 6 steps, whereas other AR models rollout for 2 steps during training. For model that has complex parameters, each coloat parameter count as two paramaters.
FNO3D
F-FNO3D
Dil-ResNet
FactFormer
Model
AR
LM
AR
LM
AR
LM
AR
LM
P avg. error
0,8080
0.4634
0.3151
0.3264
0.1725
0.1778
0.2989
0.2545
p final error 1.1285
0.6522
0.4250
0.4159
0.2573
0.2448
0.4407
0.3431
u avg. error 0.3967
0.3382
0.2298
0.2303
0.1143
0.1250
0.1775
0.1670
u final error
0.6561
0.4735
0.2799
0.2850
0.1675
0.1671
0.2594
0.2218
Inf. time (s)
1.01
0.91
2.77
1.37
12.67
6.89
2.68
1.31
# params (M)
509.8
3.0
6.9
5.1
Table 2: Evaluation results of 3D isotropic turbulence. A batch size of 4 is used for inference. LM models predict 2 steps with each call to the model. Total prediction length is 10 steps.
8

---

FNO3D
F-FNO3D
Dil-ResNet
FactFormer
Model
AR
LM
AR
LM
AR
LM
AR
LM
avg. emror
0.1607
0.1344
0.1038
0.1236
0.0843
0.0999
0.1017
0.0942
final error
0.1775
0.1287
0.1415
0.1219
0.1070
0.1062
0.1693
0.0941
avg. error 0.5198
0.4255
0.3419
0.3713
0.2378
0.2747
0.3537
0.2592
u final error
1.0245
0.6718
0.8655
0.6146
0.5372
0.5023
0.7881
0.4482
Inf. time (s)
3.19
1.47
6.35
2.75
27.49
6.94
5.61
2.62
params (M)
509.8
3.0
6.9
4.6
Table Evaluation results 3D smoke buoyancy. A batch size is used for inference. LM models predict 4 steps with each call to the model. Total prediction length is 16 steps.
Relative E2 norm (x10-2)
4.50
FNO
Linear
F-FNO
Relative Cy norm
Factorized
2.50
Galerkin Transformer
DilResNet
1.50
1.00
0.75
0.2
0.8
1.0
0.50
141
211 Resolution
421
Prediction (s)
Figure 3: Error on 2D Darcy flow with different training resolutions. Galerkin Transformer's result is taken from the original paper [16].
Figure 4: Error trend of different Trans- former models 2D Kolmogorov flow with 128 x 128 grid.
Fwd.
Fwd. Bwd.
Model
Avg. rel. L Norm
Enc. time (s)
Prop. time (s)
Time (s)
Mem. (MB)
Factorized attention
0.1529
0.0202
0.0954
5217
Linear attention (matmul) Linear attention (einsum)
0.1370
0.0112
0,3013
12029
0.5853
0.1333
0.2938
12029
Table 4:
Comparison between factorized and linear attention on their forward/backkward computational cost, Mem. denotes the peak memory usage. The benchmark is carried out using PyTorch 1.8.2 on RTX 3090,
with batch size of 4. Enc. time: the time spent obtaining the latent encoding, primarily includes attention layers and foedforward layers after attention layer; Prop. time: the time used propagate dynamics in the latent space with a a 3-layer MLP.
4.3 Comparison against full attention
In this subsection, will present ablation study of the proposed factorized attention mechanism with softmax-free attention (denoted as "linear attention") previously applied to PDE modeling 16, 67]. More specifically, employ the attention from Li et al. [67] (in the form of (4)) to replace factorized attention in (10), with K, V normalized column-wise via instance normalization, e.g. ||V. 1. To accommodate for the memory cost of linear attention, we further downsample the 2D Kolmogorov flow discussed in the last subsection to 128 x 128 grid and train both linear and factorized attention models on it (with latent marching and pushforward trick).
Comparison of performance We compare the accuracy and computational cost the two attention mechanisms in Figure 4 and Table 4. While in principle full attention could have better approximation capacity than factorized attention, practice we find that it performs worse than factorized attention on this problem studied. Specifically, its rollout is less stable and results in degraded accuracy. We hypothesize that this is due to the instability of iteratively calculating the attention matrix of large size, rolling out the prediction requires recursively calling the model multiple times. In

---

addition to the accuracy improvement, the benchmark on computation empirically demonstrates the computational efficiency improvement of factorized attention over linear attention. We provide more detailed comparison between factorized attention and linear attention in the Section D of Appendix, where we observed consistent efficiency improvement with different grid sizes and model sizes.
Pattern of attention matrices We also investigate the structure of different attention matrices. By construction, when using softmax-free attention to compute the kernel integral in (2), the kernel matrix A QK" is going have low-rank structure since rank (A) min((ank rank (K)), Q, E I xd and usually.V >> d. After training, we compute the attention matrices based on 100 samples and conduct singular value decomposition (SVD) on them. We define the total energy of the spectrum the of singular values E Σ; Ji, where o; is the i-th singular value and report the normalized cumulative energy histogram 6y = Σ oi L; Jj in Figure Sa, 5b, Sc. For each layer, the histogram is averaged across the attention matrices of all heads. We observe that for linear attention, its rank is relatively low as less than 5% the singular values capture over 90% of the total energy, which is similar to the trend observed from previous works studying the rank of standard softmax-attention [6, 26, 120J. Note that the spectrum of linear attention is based on a truncated SVD and therefore its rank will be the full attention matrix hints the potential to approximate with or decomposed into smaller and more compact matrices, and our proposed factorized scheme is one example.
lower if full SVD performed. The highly low-rank structure of
0.
Comulritve ere gy
0.
Cumulatbv layor
0.4
8
8.
First layer
Last layer
0.2
layer
6.
117
Fraztion total siigular values
Fraction
singular valuess
Roction total singular volues
(a) Full attention matrix A
(b) Attention matrix A(1)
(c) Attention matrix A
Figure 5: Spectrum of different attention matrices. (a): Attention matrices from softmax-free attention. (b), (c): Attention matrices for axis x and y from FactFormer. The vertical line indicates the fraction of singular values that capture 90% of the total energy. Since the full attention matrix has relatively large size (16384 16384). Therefore its spectrum is computed via TruncatedSVD [39] with top 1024 components truncated.
Conclusion
In this work, we propose an end-to-end Transformer for PDE modeling, which features a learnable projection operator and a factorized kernel integral. We demonstrate that the proposed model balances efficiency and accuracy well, making it a promising and scalable solution for PDE surrogate modeling. However, the proposed attention mechanism still not free from the curse of dimensionality. The computation of the factorized kernel integral requires evaluating the function all S₁ x SS x... S.. grid points. A future direction could be extending the factorization scheme to a more efficient tensor decomposition format like tensor-tras train. The proposed model currently exploits the uniform structure of the underlying grids and use mean pooling when doing projection, but non-une un-fonifom quadrature weight will be necessary when applying non-uniform grids. It is also observed that the proposed model and other neural PDE solvers can be unstable due to the error accumulation when solving time-dependent systems.
Acknowledgement
This work is supported by the National Science Foundation under Grant No. 1953222. The authors would like to thank the annonymous reviewers and area chair for their efforts and valuable feedhack during the reviewing process. The authors would like to thank Dr. Shuhao Cao from University of Missouri Kansas City for the comments regarding the backpropagation of feedforward layer and attention layer in Transformer. The authors would also like to thank Zhi Ye for the suggestions benchmarking the computational cost of the model.
10

---

References
[1] Large-Scale Dynamics and Transition to Turbulence in the Two-Dimensional Kolmogorov URL https: //arr//aca/ccc aiaa.
Flow, pages 374-396. doi: 10. org/doi htg/doi/abs..e aas///ab//// 10. 2514/5 9781608866111111118 0ww.:/.........00 0396.
[2] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2014. htRps/ts::/aare wttaaaiiiiiiiaaaaaarrreerrreeeee awrggabs/1499...
[3] Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P Brenner. Learning data-driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(31):15344-1533 2019.
[4] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document trans- former, 2020.
[5] Kaushik Bhattacharya, Bamdad Hosseini, Nikola B. Kovachki, and Andrew M. Stuart. Model reduction and neural networks for parametric pdes, 2021.
[6] Srinadh Bhojanapalli, Ayan Chakrabarti, Himanshu Jain, Sanjiv Kumar, Michal Lukasik, and Andreas Veit. Eigen analysis of self-attention and its reconstruction from partial computation, 2021.
[7] Nicolas Boullé, Seick Kim, Tianyi Shi, and Alex Townsend. Learning green's functions associated with dime-dependent partial differential equations. Journal of Machine Learning Research, 23(218):1-34, 2022.
[8] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for pde modeling. arXiv preprint arXiv:2209 04494, 2022.
[9] Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data aug- mentation for neural PDE solvers. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Con- ference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 2241-2256. PMLR, 17-23 Jul 2022. URL pattps://pron htttp:://proceedings mlr, .arrss////////. v162/brandstetterr22 html.
[10] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint arXiv:2202.033770 2022.
[1I] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K. Gupta. Clifford neural layers for pde modeling. 2023.
[12] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers, 2023.
[13] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhari- wal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
[14] Salva Rühling Cachay, Peetak Mitra, Haruki Hirasawa, Sookyung Kim, Subhashis Hazarika, Dipti Hingmire, Phil Rasch, Hansi Singh, and Kalai Ramea. Climformer-a spherical trans- former model for long-term climate projections.
[15] Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis. Physics-informed neural networks (pinns) for fluid mechanies: A review. Acta Mechanica Sinica. 37(12):1777-1738, 2021.
[16] Shuhao Cao. Choose a transformer: Fourier galerkin. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances Neural In- formation Processing Systems, volume 34, pages 24924-24940. Curran Associates, Inc., 2021. URL http:::/ htttpp:://proceedingss wwwpeediiiinnnnnniiineeeeeeeeeee weurips.cc.//eeoooe cttp.ccc.pppr.r..c paper fiw.filleepeeee wwtp:/fffeeeeeeeeeeeeeee plles/paperr2221/fiiee 00011444ees11990225055444111.113 Pappeaapprppppeee pdf.

---

[17] Shuhao Cao, Peng Xu, and David A. Clifton. How to understand masked autoencoders, 2022.
[18] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers, 2020.
[19] Gary J. Chandler and Rich R. Kerswell. Invariant recurrent solutions embedded in a turbulent two-dimensional kolmogorov flow. Journal of Fluid Mechanics, 722:554 595, 2013. doi: 10. 1001017fmm.2233.22..
[20] Ashesh Chattopadhyay, Mustafa Mustafa, Pedram Hassanzadeh, and Karthik Kashinath, Deep spatial transformers for autoregressive data-driven forecasting of geophysical turbulence. In Proceedings of the 10th international conference on climate informatics, pages 106-112, 2020.
[21] Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions on Neural Networks, 6(4):911-917, 1995. doi: 10.1 10.110977.3922553.
[22] Rewon Child, Scott Gray, Alec Radford, and Ily Sutskever. Generating long sequences with sparse transformers, 2019.
[23] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2022.
[24] Mengyu Chu and Nils Thuerey. Data-driven symthesis of smoke flows with CNN-based feature descriptors. ACM Transactions on Graphics, 36(4):1-14, jul 2017. doi: 10.1145/3072959 3073643. URL https: /ttttp:://doi..r//1/0/0 org/10. 1115111333229555 200 3073643.
[25] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/810.00805. 2018.
[26] Yihe Dong, lean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth, 2021.
[27] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. /929, 2020.
[28] Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zepeda-Núñez, Jamic A. Smith, Michael P. Brenner, and Stephan Hoyer. Learning to correct spectral methods for simulating turbulent flows, 2022.
[29] Antonio H de O Fonseca, Emanuele Zappala, Josue Ortega Caro, and David van Dijk. Contin- uous spatiotemporal transformers. arXiv preprint arXiv:2301.13338 2023.
[30] Kai Fukami, Koji Fukagata, and Kunihiko Taira. Super-resolution reconstruction of turbulent flows with machine learning. Journal of Fluid Mechanics, 870:106 170:106-110,,,00 2019. doi: 10.1017/ jf1..019.228.
[31] Zhihan Gao, Xingjian Shi, Hao Wang, Zhu, Yuyang Bernie Wang, Mu Li, and Dit- Yan Yeung. Earthformer: Exploring space-time transformers for earth system forecasting, Advances in Neural Information Processing Systems, 35:25390-25400 2022.
[32] Nicholas Geneva and Nicholas Zabaras. Transformers for modeling physical systems. Neural Nenworks, 146:272-22 289,,,2999,.. 2022.
[33] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines, 2014. URL https: haabb///441..00..
[34] John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.1353 2021.
12

---

[35] Ruchi Guo, Shuhao Cao, and Long Chen. Transformer meets boundary value inverse problems. arXiv preprint arX1v:2209.1497,, 2022.
[36] Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet- operator learning for differential equations, 2021.
[37] Jayesh K. Gupta and Johannes Brandstetter. Towards multi-spaatooeeppralassaal generalized pde modeling, 2022.
[38] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spaaiotempporal-caad generalized pde modeling, arXiv preprint 1rxXiv:2209.16616, 2022.
[39] Nathan Halko, Per-Gunnar Martinsson, and Joel A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions, 2010.
[40] Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34): 8505-8510, 2018.
[41] Xu Han, Han Gao, Tobias Pfaff, Jian Kun Wang, and Li-Ping Liu. Predicting physics in mesh-reduced space with temporal attention. arXiv preprint arXiv:2201.09113, 2022.
[42] Zhongkai Hao, Songming Liu, Yichi Zhang, Chengyang Ying, Yao Feng, Hang Su, and Jun Zhu. thysics-informe machine learning: survey on problems, methods and applications. 2023.
[43] Zhongkai Hao, Chengyang Ying, Zhengyi Wang, Hang Su, Yinpeng Dong, Songming Liu, Ze Cheng, Jun Zhu, and Jian Song. Gnot: A general neural operator transformer for operator learning. arXiv preprint arXiv:2302. 14376, 2023.
[44] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
[45] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers, 2020. URL httppp/////ppperreeeeeeee///eeee H1e5.JBtDr.
[46] Philipp Holl, Vladlen Koltun, and Kiwon Um. phiflow: A differentiable pde solving framework for deep learning via physical simulations.
[47] Ameya D Jagtap and George Em Karniadakis. Extended physics-informed neural networks (xpinns): A generalized space-time time domain decomposition based deep learning framework for nonlinear partial differential equations. Communications in Computational Physics, 28(5): 2002-2041 2020.
[48] Steeven JANNY, Aurélien Bénéteau, Madiha Nadri, Julie Digne, Nicolas THOME, and Christian Wolf. EAGLE: Large-scale learning of turbulent fluid dynamics with mesh trans- formers. In The Eleventh International Conference on Learning Representations, 2023. URL
https::/opeereeviee www/ fttt/fommmmdmfffffm IXADmmmSS
[49] Chiyu "Max" Jiang, Soheil Esmaeilzadeh, Kamyar Azizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A. Tchelepi, Philip Marcus, Mr Prabhat, and Anima Anandkumar. Meshfreeflownet: A physics-constrained deep continuous space-time super-resolution frame- work. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2020, doi: 00.1099/SC41444522000000000
[50] Pengzhan Jin, Shuai Meng, and Lu Lu. Mionet: Learning multiple-input operators via tensor product. SIAM Journal on Scientific Computing, 30000)3399-A35144 2022.
[51] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596777))33377))3355588
2021.
13

---

[52] George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed informed machine learning, Nature Reviews Physies, 3(6):422 440, 2021.
[53] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5156-5165. PMLR, 13-18 Jul 2020. URL https::////.... http///pooeeeeeeee//eee mlr press/ / 19/ wtt.///aaaaaraaosaaaaaa0 html.
[54] Georgios Kissas, Jacob Seidman, Leonardo Ferreira Guilhoto, Victor M. Preciado, George J. Pappas, and Paris Perdikaris. Learning operators with coupled attention, 2022.
[55] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URL https: A //pp//eeeeneeeeeeeee
[56] Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan computational fluid dynamics. Proceedings of the National Academy of Sciences, 1 118(21), may 2021. doi: 10. (.073/pnas.210178411 URL
Hoyer. Machine
https://doi.
org/ 1ww.//pppaa........ 2101784118.
[57] Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Rev., 5155):555,000, aug 2009. ISSN 0036-1445. doi: 10.1137/0707011XX URL https: //doi.org// 10. 11..111177770011111
[58] Jean Kossaifi, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high resolution PDEs, 2023. URL https:
/ //openreiies wwwpriiwwweeeeeeeeee fft://oormmmmm/ wewpe///oomiiddoooooooo Tas=ppo-oRsst44
[59] Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators, 2021.
[60] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Arizzadeneshelii Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint (xXiv:2108.0848, 2021.
[61] Remi Lam, Alvaro Sanchez-Gonzales Matthew Willson, Peter Wirnsherger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Faton-Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Jacklynn Stott, Oriol Vinyals, Shakir Mohamed, and Peter Battaglia. Graphcast: Learning skillful medium-range global weather forecasting, 2022.
[62] G. Lamorgese, D. A. Caughey, and S. B. Pope. Direct numerical simulation of homogeneous turbulence with hyperviscosity. Physics of Fluids, 17(1). 12 2004. ISSN 1070-6631. doi: 10.1066/1.1833. URL https: /tw.///oooioa//////ooo or.//o. 1063/1 1833415. 015106.
[63] Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition, 2015.
[64] Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids, 2019.
[65] Zijie Li and Amir Barati Farimani. Graph neural aetwork-accelrratee lagrangian fluid simulation. Computers & Graphics, 103:201-211, 2022. ISSN 0097-8493. doi: https: //doi //doi.or//10.1 wwwp://...../.............
URL https://tw/www sciencedirect. com/science/
arricleepp1//00978488220000
[66] Zijie Li, Tianqin Li, and Amir Barati Farimani. TPU-GAN: Learning temporal coherence from dynamic point cloud sequences. International Conference on Learning Representations, 2022. URL https:t/openrreeiieee www.:ww..eeeeeeee net 19FFk.
[67] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https:/ http:///pppeeeeee.eemme wttp://www..eeetoommddd/dddeeeea
14

---

[68] Zongyi Li. Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08892 2020.
[69] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations, 2020.
[70] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations, 2020.
[71] Zongyi Li. Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.034833 2020.
[72] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations, 2023.
[73] Xinliang Liu, Bo Xu, and Lei Zhang. Ht-net: Hierarchical transformer based operator learning model for multiscale pdes. arXiv preprint arXiv:2210. 10890, 2022.
[74] Xinliang Liu, Bo Xu, and Lei Zhang. Mitigating spectral bias for the multiscale operator learning with hierarchical attention, 2023.
[75] Winfried Lötzsch, Simon Ohler, and Johannes Otterbach. Learning the solution operator of boundary value problems using graph neural networks. In ICML 2022 2nd Al for Science Workshop, 2022. URL https: http:::////ppereeeee /pwp//pppeeeeeeeeeoe met/forumnid
[76] Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv: 1910.0319,,, 2019.
[77] Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis. Deepxde: A deep learning library for solving differential equations. SIAM review, 63(1):208-228, 2021.
[78] Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based on FAIR data. Computer Methods in Applied Mechanics and Engineering, 393:114778, apr 2022. doi: //doi. org. 10. 1016 wttpr///..0006/..... cma. 2022. 114778.
wtt.j/.....22.11177.7.. URL https:
[79] Minh-Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-baan based neural machine translation, 2015. URL https::/arxiv... /rg////////...... 04025.
[80] Xindian Ma, Peng Zhang, Shuai Zhang, Nan Duan, Yuexian Hou, Dawei Song, and Ming Zhou. A tensorized transformer for language modeling, 2019,
[81] Mikael Mortensen and Hans Petter Langtangen. High performance python for direct numerical simulations of turbulent flows. Computer Physics Communications, 203:53-65, jun 2016. doi: 10. (1001//.cpc.2216..22.
URL https: //oii org/ 2org/10.00 1a....11122FFF cpc 2016. 02.00...00 005.
[82] Tan Nguyen, Minh Pham. Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. Fouri- erformer: Transformer meets generalized fourier integral theorem. Advances in Neural Information Processing Systems, 35:29319-299 293355, 2022.
[83] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. Climax: A foundation model for weather and climate, 2023.
[84] Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov. Tensorizing neural networks, 2015.
15

---

[85] I. V. Oseledets. 33(5):2295-2317, 2011.
Tensor-train decomposition. SIAM Journal on Scientific Computing.
doi: 10.1137/0907522286 URL https://dot org 10. 1137/
090752286.
[86] Oded Ovadia, Adar Kahana, Panos Stinis, Eli Turkel, and George Em Karniadakis. Vito: Vision transformer-operator. arXiv preprint arXiv:2303.08891, 2023.
[87] Guofei Pang, Lu Lu, and George Em Karniadakis. fpinns: Fractional physics-informed informed neural networks. SIAM Journal on Scientific Computing, 41(4):A2603 A2626, 2019.
[88] Pranshu Pant, Ruchit Doshi, Pranav Bahl, and Amir Barati Farimani. Deep learning for reduced order modelling and efficient temporal evolution of fluid simulations. Physics of Fluids, 33 (10):107101, oct 2021. doi: 10. 02111066.0022000 URL https:///ol htgg///ooo/////////////o/// 10. 1063%225. 0062546.
[89] Jaideep Pathak, Mustafa Mustafa, Karthik Kashinath, Emmanuel Motheau, Thorsten Kurth, and Marcus Day. Using machine learning to augment coarse-grid computational fluid dynamies simulations. arXiv preprint arXiv:2010.00077 2020.
[90] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopad- hyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators, 2022.
[91] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez and Peter W. Battaglia. Learning mesh-based simulation with graph networks, 2021.
[92] Lukas Prantl, Benjamin Ummenhofer, Vladlen Koltun, and Nils Thuerey, Guaranteed conser- vation of momentum for learning particle-based fluid dynamies, 2022.
[93] Alfio Quarteroni and Alberto Valli. Domain decomposition methods for partial differential equations. Number BOOK. Oxford University Press, 1999.
[94] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In J. Platt, D. Koller, Y. Singer, and S. Roweis, editors, Advances in Neural Information Processing Sys- tems, volume 20. Curran Associates, Inc., 2007. URL http:ttttt//ppppeedeiiiii neurips......p........ ca/ http:///fffeatttteaeaaaaaaaaaaa .dd.
[95] Md Ashiqur Rahman, Zachary E. Ross, and Kamyar Azizzadeneshhlli U-no: U-shaped neural operators, 2023.
[96] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. mhysics-informen neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686- 2019.
[97] Stephan Rasp, Peter D. Dueben, Sebastian Scher, Jonathan A. Weyn, Soukayna Mouatadid, and Nils Thuerey. WeatherBench: A benchmark data for data-driven weather forecasting. Jour- nal of Advances in Modeling Earth Systems, 12(11), nov 2020. doi: 10. 10.1029/2020ss002200. URL https://doi. http:::////i.og///////// 10. 1029k222002mm00022000.
[98]
Robert S. Rogallo. Numerical experiments in homogeneous turbulence. 1981.
[99] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models, 2022.
[100] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physies with graph networks, 2020.
[101] Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities, 2020.
[102] Dule Shu, Zijie Li, and Amir Barati Farimani. mhhysies-infomeet diffusion model for high-fidelity flow field reconstruction. Journal of Computational Physics, 478:11 2023,
2023.
16

---

[103] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittieserr, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):
484 484, 2016.
[104] Kimberly Stachenfeld, Drummond B. Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanzhez-Gonzalez. Learned coarse models for efficient turbulence simulation, 2022.
[105] Jianlin Su. Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2022.
[106] Luning Sun, Han Gao, Shaowu Pan, and Jian-Xun Wang. Surrogate modeling for fluid flows based on physics-constraing deep learning without simulation data. Computer Methods in Applied Mechanics and Engineering, 361:11272 2020.
[107] Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains, 2020.
[108) Jingwei Tang, Vinicius C Azevedo, Guillaume Cordonnier, and Barbara Solenthaler. Neural green's function for laplacian systems. Computers & Graphics, 107:186-196, 2022.
[109] Nils Thuerey, Konstantin Weißenow, Lukas Prantl, and Xiangyu Hu. Deep learning methods for reynolds-averaged navier-stokes simulations of airfoil flows. AIAA Journal, 58(1):25-36, 2020.
[110] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators, 2023.
[111] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Phiil-pp Morency, and Ruslan Salakhutdinov. Transformer dissection: A unified understanding of transformer's 's attention via the lens of kernel, 2019.
[112] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization, 2017.
[113] Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey. Solver-in-the- loop: Learning from differentiable physics to interact with iterative pde-solvers. Advances in Neural Information Processing Systems, 33:6111-6122 2020.
[114] Benjamin Ummenhofer, Lukas Prantl, Nils Thuerey, and Vladlen Koltun. Lagrangian fluid simulation with continuous convolutions. In International Conference on Learning Represen- tations, 2020. URL htt.tttss//opeereviie http:/iwwwww..eeeeeeeeeeeeeeeeee (tt.:/oorrmmnddB1DD011DD0
[115] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceeding" www.gg..nuris.... pctppcaae///e////1111 (ttee2245554ddee11add0333..3.... c48433cacc444aaaaaaa Paper. pdf.
[116] Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, KDD '20, page 1457-1466, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 9784499979 doi: 10.1145/3344486.4400
URL https: /doi org/10.. 1145/339444 /39543334486... 3403198.
[117] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informer deeponets. Science Advances, 7(40): eabi8605, 2021. doi: 10.1 2..112//ciiadv.ab8866 URL. https: :ttps:///wwwe //www. science. ciencc.org//oi..................
abs/10. /0. 1126/saa /st///dddvv...m abi8605.
17

---

[118] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-infommne -nffrse- deeponets. 2021.
[119] Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is you need for training physics-informed neural networks, 2022.
[120] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 2020.
[121] Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. U-fno-an enhanced fourier neural operator-based deep-learning model for multiphase flow. Advances in Water Resources, 163: 1031004 2022.
[122] Matthew A Wright and Joseph E Gonzalez.
Transformers are deep -iieite-dimessionll non-mercer binary kernel machines. arXiv preprint arXiv:2106.01506 2021.
[123] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström- based algorithm for approximating self-attention, 2021.
[124] Yinchong Yang, Denis Krompass, and Volker Tresp. Tensor-train recurrent neural networks for video classification, 2017.
[125] Fisher Yu and Vladlen Koltun. Multi-scale context aggregation by dilated convolutions, 2016.
[126] Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences, 2021.
[127] Min Zhu, Handi Zhang, Anran Jiao, George Em Karniadakis, and Lu Lu. Reliable extrapolation of deep neural operators informed by physics sparse observations. Computer Methods in Applied Mechanics and Engineering, 412:116064, 2023.
18

---

A
Model implementation details
The major hyperparameters are listed in Table
Hyperparameter
2D Kolmogorov
3D turbulence
3D smoke
2D Darcy flow
Hidden dimension
128
128
128
128
Depth
4
4
3
3
Heads
8
6
12
Kernel dimension
128
192
192
128
Input encoder
2D Conv MLP
2D Conv
2D Conv
MLP
Output decoder
MLP
MLP
MLP
Table 5: Major hyperparameters for FactFormer
Hidden dimension indicates the number of channels in the latent space. Depth denotes the number of attention layers. Before entering each attention layer, we modulate the latent encoding with positional encoding: z; f- 24 + (rr;), where x; is the Cartesian coordinate of latent encoding Pi, : + Fd is a random Fourier feature mapping [94, 107] with a learnable linear transformation. Kernel dimension is the dimension of each head, which is equivalent to dg, the number of function bases used to compute the kernel: <(z,g) Σi du (1)) We train the model with AdamW optimizer and cyclic learning rate scheduler with a maximum learning rate 3e - 4, similar to prior foreeformer-for-DEee works [16, 43, 67].
For 3D smoke buoyancy and 2D Darcy flow problem, we append a CNN block after attention layers to better account for the boundary values. The CNN block has U-shape arrangement with 4 CNN layers, with all layers using kernel size of and padding size of (pad with zeros). The first CNN layer has a stride of 2, while other layers have stride of 1. The stride-2 convolution will downsample the data by half, so nearest upsampling is applied between the second and third CNN layers to recover the spatial resolution.
On top of every model (including baselines we will discuss in the next section), we use 2D convolutional layer to compress the temporal dimension if it is time-dependent problem. Concretely, we first reshape the input into (i, where N is the number of spatial grid points and Th is the number of input frames. Then we apply 2D convolution filters of size (1, Tin) to compress the temporal dimension to At the bottom of every model, we use three-layer MIP to project the latent encoding back to variables of interest such as pressure and velocities. In addition, we adopt a curriculum training strategy for all latent marching models, where we only march for 1 step at the beginning of the training and don't do any pushforward. Then we gradually increase the latent marching steps throughout the training and apply pushforward when the model has been trained for around 6% of the total epochs.
B Baseline implementation details
In this section, we provide the full details of baseline models, namely FNO, F-FNO, Dil-ResNet, and linear attention Transformer.
For Fourier Neural Operator, the implementation is taken from Li et al. [72]'s official imple- mentation: https::/ //gtthub..tttt.tt www./ooo/omaaeee.ee httmm//eraalopratrrrpppoo/ Fourier Neural Operator the implementation is taken from https: htsps:///ithhll com/ alaadairtatr fourierflow. We add group normalization before the final fully-connected layer. We use hidden size of 96, a mode number of 12 for 3D problems, 24 for 2D turbulence, 20 for Darcy flow, and a layer number of 4.
informedd And for Factorized
For Dee-ResNet, we adopt the implementation from Gupta and Brandstetter [37]:https: //gittu.. com/microsof troooott/et t/pdearenntt
Compared to Gupta and Brandstetter [37] and Stachenfeld et al. [104], we simplify the setting by truncating the number layers inside each block, where we use dilation layers [1, 3, 8, 3, 1] (1, 2, 4, 2, is chosen for the 3D problem it has slightly better performance) inside each block instead of [1, 2, 4, 8, 4, 2, 1]. The primary reason that without truncation, the training on 3D problems will take over week (and cannot fit into single A6000 GPU for 2-step
19

---

rollout training), which is significantly slower than other models. In summary, we use 3 residual blocks (each with 5 CNN layers of width 128), an MLP-based or 2D convolution- (for convolution in temporal domain) encoder/decoder, with group normalization inserted between every residual block. The implementation of the linear attention Transformer follows OFormer's [67] attention implementation https: //gt thub. com/ Baratiiabb 2atp//Formeee ormer, with Galerkin style normalization scheme [16]. Other hyperparameters are kept the same as the hyperparameters listed in Table 5 2D Kolmogorov. For non- Transformer models, they are trained with Adam optimizer and decay learning rate from 5e 4 to
throughout the training.
All the experiments are carried out using Py Torch 1.8 except for FNO/F-FNO experiment, which uses PyTorch 1.13 for optimizing complex-valuse parameters. We train all LM models for 100k iterations and AR models for 64k iterations of gradient updates.
C Visualization of error trend
This section includes the average frame-wise error trend for the time-dependent systems we have investigated.
DilRes
FNO
FactFormer F-FNO
Relative 4 nom
0.0
0.2
0.4
0.5
0.8
1.0
Prediction time (s)
Figure 6: Error trend of vorticity w on 2D Kolmogorov flow. Dashed line: AR; Solid line: LM
0.8
DilRes
DilRes
FactFormer
FactFormer
Relative norm
0.4
norm
0.4
0.2
0.3
Prediction
Prediction
Figure 7: Error trend of pressure p on 3D isotropic turbulence.
Figure 8: Error trend of velocity 3D isotropic turbulence.
20

---

0.6
FNO
FNO
0.5
FactFormer
0.5
FactFormer
F-FNO
Relative E. n.om
0.4
Rolative C. norm
0.3
0.3
0.1
0.0
Prediction time (6)
10
12
0.0
Prediction time (5)
10
Figure 9: Error trend of marker field d on 3D smoke buoyancy.
Figure 10: Error trend of velocity field u on
3D smoke buoyancy.
D Further ablation study
This section includes further study and numerical experiments on the proposed model.
Note backpropagation Consider single (softmax- free)) attention block that consists of attention layer and two layer feedforward network (for simplicity we consider single-headed case):
U1 Au(Uo), U2 Wi)
(11)
where Uo RRx is the input, Wi, W2 RRdd are learnable weights for the feedforward network. Given loss function I(.) Nxd + IR (for example, squared error), define U₁ Wi. 01 == h, @r(Ui)) a02
at; == the gradient for the weights and input in the feedforward network are (o
denotes element-wise multiplicattin):
01
oW h
(12)
0l (hW, θ
(13)
OW
0L (NW? o g)
(14)
aU,
For linear (softmax-free) dot-product attention: U1 Au(Up) = QKTV, where Q Uo y: K U.W, V U,W, v. and thus U, = WWW,W, Ug W.,,, the gradient of weights are:
Ø1 U/²'
Øl
W UaWn,
(15)
oW, 9 OU
Øl U U, Øl TU.W, OU
(16)
OW,
01 ow,
U.U,W. of Øl
(17)
OU
where Uo E Nxd has appeared three times in each calculation and thus in backpropagation the computational cost attention layer is generally more expensive than the feedforward layer where involves more matrices that grow exponentially with respect to the spatial resolution.
Next we provide a comparison between the backpropagation of linear attention and the proposed factorized attention in scalar summation form. For simplicity, we only consider the attention layer. For linear attention, be written as:
Zi
Σ Quml m=l
K jm Vgc),
(18)
21

---

where Zic denotes the i-th row and j-th column of matrix Z and similar for other matrices, Q
XW. K = XW V = XW, and X Rxd is input. The gradient of parameters are computed
as:
N 5ie),
-xir() E
(19)
a(W g)r.
1=1
07,,
-ii
E j=1
(20)
((Werr
07..c
En
Qim (E Kam)er)
if = otherwise.
(21)
a(W.)r,
For the proposed factorized attention, can written as:
5
82
Sn
ii = 2
A 1232
(n) Jacn
Vic. j:= (Ji.ja...... In),,, i == i1, d2.......................
in
h= 1==1
Jn
(22)
where (mm) (Kmmm)J (A(m) E RR N S x X S,) is the axial kernel matrix as defined in (8), (m) E R xd is the axial projection along m-th axis defined in (6).
For W., o. gradient is computed as:
877,c
Σ E ΣS. 1
1
A Xjr
ifs=
(23)
a(W.)r.
otherwise.
For w(n), (x) their gradients computed as:
oZ,
x(()
81
82
Sa
(1)
(2)
(
An-id-- Vgck(n)
(24)
a(W )n.
lada
Jniot
==
1==
J==1
d ((W!")) (W("))r.
OZ,cc
in
5
S,
(1)
((m-1)
In-1-
VicX(())
(25)
rr.
j1=1
j==1
The major difference is that for factorized attention the summation is taken over each axis separately while for linear attention is taken over all N grid points.
discussed in Section A, the kernel dimension indicates how many function bases are used to evaluate the kernel and larger kemel dimension is beneficial to the learning capacity of the model. As shown in Figure la, 1 11b, linear attention's training cost increases more significantly than the factorized attention kernel dimension increases, since its complexity quadratic with respect to kernel dimension. Factorized attention's 's computational efficiency can be further improved by reducing the spatial resolution, leveraging techniques such learning the mapping in the latent space (similar to latent diffusion model [99]). multi-scale network architecture that resembles multigrid methods [35, 74], domain decomposition [47, 93]. Furthermore, the training cost of factorized attention is also relatively lower than lincar attention on 3D domain as shown in Figure 12a, 12b.
Runtime comparison
22

---

Factorbed
Linear
Factorized
Linear
D0M
COM
20000
13200
feed Swe.. Time (xecc
19933
13000
5324
1013
1.95
0.00
0.19
182
256
Komel dimension
Kemei dimension
(a) Forward backward peak memory usage.
(b) Forward backward time per iteration.
Figure 11: Benchmark of factorized attention and linear attention 2D domain (with 128 x 128 grid) with varying kernel dimension. Benchmark is done RTX 3090 with PyTorch 1.8.2 and batch size of Hyperparameter setting the same as Table 5-2D Kolmogorov flow. "OOM" denotes out of memory.
Factorized
Linear
Fectorized
Lineer
00M
2,00
1.75
Mamory Lsage (MB)
40000
(sec)
30000
20000
14593
13489
10000
0.00
Grid size
Grid size
(a) Forward + backward peak memory usage.
(b) Forward backward time per iteration.
Figure 12: Benchmark of factorized attention and linear attention 3D domain with varying grid size. Benchmark done on an A6000 with PyTorch 1.8.2 batch size 1. Hyperparameter setting the same as in Table 5-3D isotropic turbulence.
Model scaling performance We study the impact of number of heads and size of kernel dimension prediction loss. For each direction of hyperparameter search, we fix the value of other hyperparameters to that shown in Table 5. The ablation experiments are conducted on 2D Kolmogorov flow (sampled from 128 x 128 grid) with splitting different from the Evaluation Section in the main body of the paper. As shown in Figure 13b, the number of attention heads has crucial impact the final performance. The model's 's performance drops significantly when using fewer heads. This highlights the importance of multi-head mechanism the factorized attention. In addition, we observe that the final accuracy of our model benefits from increased kernel dimension (as shown in Figure 13a).
0.1.44
0.420
Relative loss
0.15
0.131
Relative (2 loss
0.4
0.139
0.3
0.10
0.154
0.05
0.1
0.00
0.0
Kernel dimension
Number of heads
(a) Ablation on kernel dimension.
(b) Ablation on number heads.
Figure 13: Ablation study on the key hyperparameters. Red color denotes the final choice of hyperparameter. Experiments carried out the validation fold.
23

---

Visualization of learned kernels We visualize the learned kernel as shown in Figure 14. Due to the presence of Rotary positional encoding [105], all kernels have a stationary pattern (the kernel value k(51,{) depends only on the relative distance between two points, e.g. 1,² distance: || ε2 l2). The kernel matrices also exhibits symmetric pattern despite the non-symmetric nature of dot product QK? and all kernel matrices are diagonal dominated.
134
30
20 12
(a) Layer 1, Head I: kernel of axis
(b) Layer I, Head I: kernel of axis
(c) Layer 4, Head I: kernel of axis
(d) Layer 4, Head kernel of a axis
Figure 14: Visualization of normalized attention kernel.
Influence of random seed We investigate the influence of random seeds by training the model with three different seeds. As shown in Figure 15, all models converge to the similar level of loss with marginal difference.
Seed
0.25
Seed
Seed
Relative L'2 norm
0.20
0.15
0.10
0.05
0.2
0.4
0.6
0.8
1.0
Time (s)
Figure 15: Averaged frame-wise loss trends. Each loss curve corresponds to a model initialized under specific seed.
E Dataset details
In this section we provide the details for each dataset.
2D Kolmogorov flow
The incompressible Navier-Stokes equation under vorticity form reads as,
@w(x,t) t)
1 u(x,t) Vw(x, () V2w(x, f(x),
(0,27) E (0, T],
Ø1.
Re
V T-u(x, f) 0, w(x, 0) = wo(x).
(0, (0,27)), [0,2], (26)
x (0, 2m)
where w denotes vorticity, u denotes velocity. Re denotes the Reynolds number, x = (x1, 22) denotes spatial coordinates and ʃ(-) is the forcing term that is set to f(x) - cos (n.²2)
1e.l(x (x).
The equation is periodic in all spatial directions. Compared to the cases discussed in Li et al. [72]. we set the forcing factor n to [1, 19] and introduce dragging force term 0. wx as described in
Kochkov et al. [56]. The initial condition wo is sampled from prescribed Gaussian random field

---

same as Li et al. [72]. The dataset consists of 100 trajectories for training and 20 trajectories for testing, with the length of each trajectory being 10 seconds and 160 frames.
We modify the pseudo-spectral solver (under Apache License 2.0) from https://githes thub. com/ neuraloperator/pyyyee
wtt.rr////blb/mtteeeeeeeeeeeeeee aomaasterssolerr//ee ssslerllllerrerrr /t///ooloooge wttp:/koooooooo..ommmoomoo wwwoooowwww.woooooooooooooooo.oo -p to gener- ate the data. The referenced direct numerical simulation is carried out with spatial resolution of 2048 2048 and a temporal resolution of le 4.
3D isotropic turbulence
The incompressible Navier-Stokes equation for this problem is given as:
Ōu(x, t) 0t
+ u(x,t) Vu(x, = v/ 1 P
Vp(x,t)+ + f(x),
x (0,2*) (0, T),
u. 1) = O,
x E (0, 2x) / E [0, T), (27)
u(x,0 0) = (x(),,
x E (0, 2x)²
where denotes velocity, p denotes the pressure, is the viscosity parameter, x = [21, 32, 33 denotes the spatial coordinates and f(.) is the forcing term. The equation is periodic in all three spatial dimensions. The initialization of u and the forcing settings follow Rogallo [98] and Lamorgese et al. [62] respectively, with Taylor Reynolds number set to 84 [62]. The dataset consists of 1000 trajectories for training and 100 trajectories for testing, with the length of each trajectory being 1 second and 20 frames.
We use the spectral Galerkin solver (under GNU GPL license 3.0) from bttps://githie com/ spectralDNS. The referenced simulation is carried out with a spatial resolution of 60 x 60 x 60 and temporal resolution 0.005s.
3D smoke buoyancy pressible Navier-Stokes equation (similar as above) coupled with advection equation:
The governing equations for the 3D smoke buoyancy problem are incom-
tu(x,t) t) 06
+ u(x, 8) Vu(x, t) uVu t) 1 V (), x (0, (0,7),
0d(x,t)
u(xxt) Vd(x, t) = 0,
x L), t (0,
Ot
V u(x,
E (0, L) [0, x E L)3
u(x, 0)= = 0, d(x, 0) do dx(x).
where f(x,1) 1) 10, 0, ,d((x, 1)). v) is the buoyancy factor, the velocity field u has a Dirichlet boundary
condition: uxx.) -) 0, vx E an, and the scalar density field for smoke has a Neumann boundary 0, vx E aQ. The initial condition dg (x) is random field 3 with scaling of Fourier coefficient set to 15.0, smoothness factor set to 4.0. The length of the rectangular domain L
condition: Vd(x,
set to 8. The dataset consists of 2000 trajectories for training and 200 trajectories for testing, with the length of each trajectory being 15 seconds and 20 frames.
We modify the 2D solver (under MIT license) from https: //gtthub.. http:///ooooooooo/aaeaoo.oooooo pdearee to generate the data. The solver applies an diveection-projeer scheme. The referenced simulation is carried out with spatial resolution of 64 64 64 and a temporal resolution of 0.75s.
2D Darcy flow
The equation for the 2D Darcy flow is defined as:
- (a(x) = f(z), E 1)², Mo (x) = 0, a(0, 1)?
(28)
where f(x) is the forcing function that is set constant 1. The coefficient function a(x:) is sampled from Gaussian Random Field with zero Neumann boundary condition. The data is generated via second-order finite difference solver 421 421 resolution grid. We use the pre-generated dataset from Li et al. [68] (under MIT license). The dataset consists of 1000 samples for training and 100 samples for testing.
Immleemntedd with phijlow's Noise class, seetht httpt:::::///ttpepeee -bs wwt.unnbbb..gg.bbbb.p.go. 10 So/Phi/Floo/ph//ielldd
25

---

Results visualization
In this section, we provide exemplary visualization of the model's prediction. For 3D problems, the cross-section at the middle of the first axis is shown.
T-40+0.2
10+0.4
10+0066
10+0.8
10+1.0
Prediction
Error
Figure 16: Sample of 2D Kolmogorov flow.
T=00+0,2
10+0.4
10+0.6
10+0.8
10+1.0
Prediction
Reference
Rrror
Figure 17: Sample of 2D Kolmogorov flow.
26

---

T=0.6
0.9
Prediction
Reforence
Error
Figure 18: Pressure in 3D isotropic turbulence.
T=0.6
0,9
Prediction
Reference
Error
Figure 19: z-component of velocity in 3D isotropic turbulence.
T-0.6
0.7
0.8
0.9
1.0
Reference
Error
Figure 20: y-component of velocity in 3D isotropic turbulence.
27

---

T=0.6
0.9
Prediction
Reforence
Error
Figure 21: z-component of velocity in 3D isotropic turbulence.
T-6
10.5
12.75
Prediction
120
Error
Figure 22: Smoke marker field in 3D smoke buoyancy.
T=6
10.5
12.75
Prediction
Reference
Error
Figure 23: n-component of velocity in 3D smoke buoyancy.
28

---

T
10.5
12.75
Prediction
Reforence
Error
Figure 24: y-component of velocity in 3D smoke buoyancy.
T=
10.5
12.75
Prediction
Reference
Krror
Figure 25: z-component of velocity in 3D smoke buoyancy.
5oise
Prediction
Reference
20339
30023
01001
Figure 26: Flow field of 2D Darcy flow.
29

---

G
Broader impact
The numerical simulation of PDEs is of extensive application in various fields, such as manufactur- ing, weather forecasting, and engineering design. Meanwhile, Transformer has shown promising performance on a wide range of data-driven applications including PDE modeling. Our work can help improve the stability and computational efficiency of the attention-bare Our experiments demonstrate that the proposed model serves as an efficient surrogate for numerical solvers, maintaining balance between accuracy and efficiency, and thus pushing the Pareto front of accuracy-efficiency. However, as there exists large variety of PDEs and each with very unique properties, there is no guarantee that type of data-driven model can rule all. Additionally, just like most concurrent works on neural PDE solvers, the long-term stability of the proposed model cannot be guaranteed. Therefore it is important acknowledge the limitations and potential risks associated with the application of neural PDE solvers.
PDE surrogate models,
Apart from enriching the existing architecture design choice of attention based models, our work also has the potential be combined with other neural PDE solvers design formulas (e-g.. explicitly take into account the relationship between different output variables), or common neural network architectures (e.g.. U-Net).
H
Schematic of Axial Transformer and FactFormer
FactFormer
Axial Transformer
input: Isice: [85, A, w. .])
Input: (size: w..d]
legend
transpose to
(1
batch size
h height
(xx) [BS, N.1
(x)
(sice: [aS x h, w.a)
width
[sioe: [OS, w. d])
d channel
g000
projection
emb
operator along
the kith axis
g00)
(20
K
Q.K.V size: [RS w.
omb point-wiss embedding
dot product
ORT (size:[BS h. h.w,w]
A (reto (size: [OS, A]]
X kth mode product
(size: h.w.dl d])
(size: [OS, w. w))
V X, 113 X, AR
CQKTOV
output: y [size: [8S, R w.dl)
output: y [size: [BS x h, k.).d)
Figure 27: An illustrative example of FactFormer and Axial Transformer applying to 2D input data, with some details such as positional encoding, multi-head mechanism and softmax in Axial Transformer omitted for simplicity. For Axial Transformer, column-wise attention block is shown as an example.
30


Reviews:
[{"review_id": "22HwTSaRMr", "reviewer": "Reviewer_YpQV", "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.    ", "questions": "1) Per Weakness #2 above, What is the computational cost of training FactFormer compared to Dil-ResNet and FNO?   \n\n2) 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.  \n\nIf these questions are (convincingly) addressed I am willing to significantly increase my score, potentially to \"Accept\". ", "limitations": "na", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1) The paper is generally well-written\n2) 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. \n3) 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?  \n\n2) (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.  \n \n3) (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.  \n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer YpQV for the detailed comments and suggestions on improving the paper. Here we would like to address your concerns and questions as below.\n\n---\n> *The motivation for creating new datasets.*\n\nOne of the major goals of our work is to improve the scalability of attention-based PDE learning model on higher-dimensional PDE problems and many existing benchmarks consider only 1D or 2D problems.\nTo this end, we created a dataset with two challenging 3D problems and a high-resolution 2D problem to stress-test our newly proposed model and existing neural PDE solvers.\n\n---\n> *Why don't the authors use the data/benchmarks from [84]?*\n\nTo our knowledge, the data and code of [84] are not publicly available yet.\n\n---\n> *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\n>  (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,\n>  and compare Factformer's performance to that reported for the baselines on that prior study.*\n\nWe appreciate reviewer's suggestion of a fair and reliable benchmark. To address your concern, we have added new experiments of FactFormer \non commonly used benchmark datasets including 2D Navier-Stokes and 2D Darcy flow from FNO (*please refer to **Table 1** in the newly uploaded one-page PDF*). We observe FactFormer also has competitive performance on these datasets.\nIn addition, we plan to open-source the code and datasets proposed in this work in the future to facilitate the development of other models in this area.\n\n---\n> *Indeed, the Dil-ResNet often outperforms the FactFormer, and therefore there seems to be little, if any, performance advantage against recent PDE solvers.*\n\nWe agree with reviewer that our proposed model does not consistently outperform Dil-ResNet. However, we do want to point out that as there exists a large variety of PDEs, there is no guarantee that one type of model can rule them all.\nPrior studies have reported that Dil-ResNet excels on 3D turbulence [1] and 2D smoke buoyancy [2]. In our newly added experiments, we show that our proposed model actually has better accuracy over Dil-ResNet on some other 2D problems.\nIn addition, unlike FNO and FactFormer, without changing the architecture, Dil-ResNet's performance deteriorates significantly when the resolution increases (**Table 1 - Darcy flow 2D** in the newly uploaded PDF). In this case, our proposed model has roughly the same level of accuracy under varying resolutions.\n\n[1] Learned Simulators for Turbulence, ICLR 2023\n\n[2] Towards Multi-spatiotemporal-scale Generalized PDE Modeling, 2023\n\n---\n> *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?)*\n\nThe problem tested is $128 \\times 128$ Kolmogorov flow (*mentioned in line 331 of manuscript, we will make this clearer based on reviewer's feedback*). On top of that, we would like to invite reviewer to also check out **Figure 1** and **Figure 2** in the **Supplementary material** , which includes runtime reports under different model sizes and mesh sizes/dimensionalities.\n\n---\n> *What is the computational cost of training FactFormer compared to Dil-ResNet and FNO?*\n\nWe have added the full runtime (forward+backward) in **Table 5** of the newly uploaded pdf. The relative order is similar to the inference time we reported in the manuscript."}, {"review_id": "SXCY1rj1tK", "reviewer": "Reviewer_hvcj", "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.", "questions": "1. How is the error trend plot computed exactly? I thought it is accumulative error but the error is not monotonically increasing. \n2. Can you also compare the model efficiency with tensorized FNO, which leverages the tensor factorization to improve efficiency of the original FNO? \n3. 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?", "limitations": "Listed in weaknesses and questions sections.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Improving the efficiency of transformer is important to the development of transformers in PDE modeling. \n2. 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. \n2. 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. \n3. 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.  \n4. 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? \n5. Overall, the accuracy of FactFormer is not consistently better than Dil-ResNet. \n\nMinors:\n1. Line 179: $L_N$ should be $L_n$. \n2. Line 302: \"time cost for ...\" -> \"time cost of ...\"\n3. Line 315: \"time cost by ...\" -> \"time cost of ...\"\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer hvcj for the insightful comments and suggestions on the paper. Here we would like to address your concerns and questions as below.\n\n---\n>*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.\n> 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.*\n\nThe projection part exploits the low-rank structure and drops information of other axis, but the value branch in the attention layer preserves the information of all dimensions and then a pointwise feed-forward network is applied to the output of the attention layer (somewhat similar to the local branch in FNO).\n\n---\n> *The authors hypothesize that it is due to the instability of rollout. An error trend plot of full attention and factorized attention is needed to verify this hypothesis.*\n\nWe've added a figure of the average rollout error trend of Linear/Factorized attention in the newly uploaded PDF (**Figure 1**).\n\n---\n> *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 allow FNO to model high-frequency patterns [1]...*\n\nWe appreciate the reviewer for the valuable reference. We will revise our description of the high-frequency bias hypothesis and incorporate the reference in the next version of the paper. We agree that theoretically FNO has the capacity to capture high-frequency patterns, but in practice, this also depends on the model's training dynamics and the problems it is applied to. [1] poses that learning the linear transform of Fourier coefficients can potentially suffer from spectral bias [2] and choose to preserve lower frequency modes during training.\n\n[1] Incremental Fourier Neural Operator, 2022. \n\n[2] On the spectral bias of neural networks, 2019.\n\n---\n> *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?*\n\nThe major reason is that the number of parameters of FNO grows exponentially with respect to the problem dimension while FactFormer does not (the projection operator of FactFormer grows linearly). The parameter count of the kernel in each FNO kernel integral layer is $O(M^nd^2)$, where $M$ is the number of truncated modes, $n$ is the number of the problem dimension. For example, using 16 modes and a hidden dimension of 128 in a FNO3D layer will result in parameter count of $O(16^3 \\times 128^2)$. We also list the FLOPs (on 3D grid) measured by DeepSpeed library below:\n\n|Model|Tensorized-FNO|Factorized-FNO|FNO|Dil-ResNet|Linear Transformer|FactFormer|\n|-----|--------------|--------------|---|----------|-----------|-------------------|\n|GFLOPs |33 | 101 | 33 |3501|1685|596|\n\n---\n> *How is the error trend plot computed exactly? I thought it is accumulative error but the error is not monotonically increasing.*\n\nFrame-wise error is shown in the error trend plot. LM models predict multiple future steps within one model call and thus result in the error not always monotonically increasing within a prediction window.\n\n---\n> *Can you also compare the model efficiency with tensorized FNO...*\n\nIn the newly uploaded one-page PDF, we added tensorized FNO to the runtime comparison (**Table 5**) and we also tested it on 2D problem (**Table 2**).\n\n\n---\n> *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?*\n\nThe low-rank property of large-sized full attention matrix $A$ hints that it is possible to simplify the original kernel integral computation without too much loss of information. \nWe approach this simplification by replacing the original kernel integral with an axial factorized integral which only involves a group of much smaller attention matrices $A^1, ..., A^n$ (with higher rank), which improves the computational efficiency. Based on the reviewer's feedback, we will revise the discussion in section 4.3 to improve the clarity.\n\n---\n> *Spotted typos...*\n\nWe thank the reviewer for the catch. We will correct them in the next version of the manuscript."}, {"review_id": "3klQhVS1l0", "reviewer": "Reviewer_UCqx", "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.", "questions": "1. is there any non-linearity or activation function used in the architecture?\n2. How well does the model perform in the zero-shot super-resolution task compared to the FNO baseline?", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is effective and scalable, making it suitable for high-resolution PDE modeling.\n2. 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer UCqx for the helpful comments and suggestions on improving the paper. Here we would like to address your concerns and questions as below.\n\n---\n> *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,\n> 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.*\n\nBased on reviewer's suggestion, we have added experiments on datasets from FNO with varying viscosities (please refer to **Table 1** in the newly uploaded PDF). We observe that FactFormer also has competitive performance on these datasets. In addition, the 2D Kolmogorov dataset in our work also has a relatively low viscosity (with Re=1000).\n\n---\n> is there any non-linearity or activation function used in the architecture?\n\nYes, as shown in Figure 2 of the main manuscript, we apply a pointwise feedforward network to the output of each attention layer, which uses GELU activation function.\n\n---\n> How well does the model perform in the zero-shot super-resolution task compared to the FNO baseline?\n\nThe proposed model can also generalize to unseen resolutions like FNO. Below we provide an experiment on Darcy flow:\n\n(Metric is relative L2 error)\n|Resolution|211(train)|421|\n|----------|----------|---|\n|FNO2D|0.0073|0.0141|\n|FactFormer|0.0058|0.0133|"}, {"review_id": "2lxQ7oHAA1", "reviewer": "Reviewer_y1WA", "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.\n\nI 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.", "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?\n- In Section 4.2, the FNO models have significantly higher parameter counts. Is there a specific reason for this?\n", "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).", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 2, "contribution": 3, "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.\n- Dramatic savings versus even the linear attention.\n- 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.\n- 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.\n- 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.\n- 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, \n- If only the author could add a more informative diagram comparing the difference between the factorized approach vs Axial Transformer.\n\n### Misc small typos\n- line 144: there should be space between RoPE and its reference.\n- 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)\".\n- line 154: \"learnable projection\" -> \"learnable projections\".\n- line 209: in (9), $Z$ should be set to $\\operatorname{Att}(U)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer y1WA for the helpful comments and greatly appreciate your recognition on our work. Here we would like to address your concerns and questions as below.\n\n\n---\n> *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,\n>  but compares that in (2), which is for a 1-D problem.*\n> \n> *The presentation in laying out the factorization could use some improvement for people more familiar with linear algebra than differential equations...*\n\nWe appreciate reviewer's suggestion on making this paper more accessible to the broader audience. In (2), we are trying to define a general domain $\\Omega$ without specifying its dimension, it can also be greater than $1$-D. In the next version of the manuscript, we will make this point clearer, and incorporate a more concrete expression explaining the difference between the standard attention and axial factorized ones, that is our proposed scheme amounts to replacing $z=AV$ with tensor-matrix product $z=V\\times_1 A^1 ... \\times_n A^n$ where $A^1, ..., A^n$ are much smaller matrices than full attention matrix $A$. \n\n\n---\n> *If only the author could add a more informative diagram comparing the difference between the factorized approach vs Axial Transformer...*\n\nIn the newly uploaded PDF (**Figure 2**), we added a diagram that describes the differences between how the proposed approach and Axial Transformer deal with 2D input.\nWe will also add this to the Appendix of the paper in the future version.\n\n---\n> *Misc small typos...*\n\nThanks for the catch, we will correct them in the next version of the manuscript.\n\n---\n> *In the comparison of SVD for attention matrices, what is the value of k? What specifically are A^{(1)} and A^{(2)}?*\n\n$k$ is the index of the singular value (used to compute the faction of singular values in the figure, i.e. k divided by total number of singular values).\n$A^{(1)}, A^{(2)}$ are sub-attention matrices corresponding to two axes $y$ and $x$, respectively.\n\n---\n> *In Section 4.2, the FNO models have significantly higher parameter counts. Is there a specific reason for this?*\n\nThis is because the number of parameters of FNO grows exponentially with respect to the problem dimension. The parameter count of the kernel in each FNO kernel integral layer is $O(M^nd^2)$, where $M$ is the number of truncated modes, $n$ is the number of the problem dimension. For example, using 16 modes and a hidden dimension of 128 in a FNO3D layer will result in parameter count of $O(16^3 \\times 128^2)$."}, {"review_id": "ePqvrrGISE", "reviewer": "Reviewer_EAHP", "summary": "This paper uses separable axis attention to reduce the complexity from being exponential in spatial dimension to linear.", "questions": "Table-4 only reports the encoding or attention layer runtime. What is the overall runtime including MLP layers?", "limitations": "NA", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\nBy 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer EAHP for the efforts spend on reviewing our work. However, we believe there are some misunderstanding and misinterpretation about our work and we would like to\nprovide some clarifications below.\n\n---\n> *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.*\n\nWe believe the reviewer has misunderstood the goal of the proposed model. The goal of the proposed approach is to alleviate the curse-of-dimensionality of applying attention to higher-dimensional PDE problems and our approach is not completely free from the exponential complexity in spatial dimension (we actually acknowledge that our model requires evaluating function value at all $N$ grid points in the conclusion section). We achieve this goal by exploiting the low-rank structure of the kernel matrix $A=QK^T$, replacing it with the product of a set of much smaller axial kernels. Through extensive experiments, we have showcased that our proposed approach enjoys notably better computing efficiency and numerical accuracy than the model using softmax-free attention.\n\nIn addition, despite the asymptotic bound of MLP being $O(Nd^2)$, it is not the major bottleneck of attention-based PDE solvers in practice. For example, on the $128 \\times 128$ problem, MLPs only account for roughly 16\\% of the total calculation time for FactFormer, and less than 5\\% for Linear Transformer. \n\n---\n> *Table-4 only reports the encoding or attention layer runtime. What is the overall runtime including MLP layers?*\n\nWe actually include most MLPs' runtime in Table-4 and we will make this point clearer in the future version of the manuscript. The only MLP that has been excluded is the one that's used for propagating dynamics in the latent space. More specifically, the model has 4 attention layers followed by 4 MLPs, and an additional propagating MLP. The *Enc. time* corresponds to the total runtime of  4 attention layers + 4 MLPs. The *Prop. time* corresponds to calling the propagating MLP 4 times to propagate system state from $z_{t}$ to $z_{t+4}$.\n\n\n---\n> *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.*\n\nThe experiment here is to provide heuristic motivation for replacing the full-attention with axial factorized attention. When parametrizing the learnable kernel integral with vanilla attention, the kernel\nmatrix $A=QK^T$ has very low rank by design (since $rank(AB)\\leq \\min[rank(A), rank(B)]$ and as reviewer points out, $rank(Q), rank(K)$ is upper bounded by channel number $d$). This motivates us to propose a axial-factorized kernel integration scheme and replace $A$ with a set of much smaller (but higher-rank) matrices $A^1, A^2, ..., A^n$. And the experiment results confirm that $A^1, A^2, ..., A^n$ indeed exhibit higher-rank structures after training.\n\nBased on reviewer's feedback, we will revise the description here to avoid ambiguous impression. Essentially, the proposed approach doesn't differ from vanilla linear attention in terms of how to calculate attention score (they are both dot product attention), the difference is how to parameterize and compute the kernel in the learnable kernel integral transform with attention. "}, {"review_id": "8SS1oP5Bi6", "reviewer": "Reviewer_TtUN", "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.\n", "questions": "None.", "limitations": "None.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The motivation for the research is clear, with the Introduction section being easy to read and understand.\n2. 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.\n3. 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. \n2. 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.\n3. 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?\n4. 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.\n5. 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.\n\nRef:\n1. Factorized Fourier Neural Operators (https://arxiv.org/abs/2111.13802)\n2. Multi-Grid Tensorized Fourier Neural Operator for High Resolution PDEs (https://openreview.net/pdf?id=po-oqRst4Xm)\n3.  HT-Net: Hierarchical Transformer based Operator Learning Model for Multiscale PDEs(https://openreview.net/forum?id=UY5zS0OsK2e)\n4. Multiwavelet-based Operator Learning for Differential Equations (https://arxiv.org/abs/2109.13459)\n5. U-FNO -- An enhanced Fourier neural operator-based deep-learning model for multiphase flow (https://arxiv.org/abs/2109.03697)\n6. Physics-Informed Machine Learning: A Survey on Problems, Methods and Applications (https://arxiv.org/abs/2211.08064)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer TtUN for the very detailed comments and suggestions. We also appreciate the recognition on our work.\nHere we would like to address your concerns and questions as below.\n\n---\n> *The experimental section is somewhat weak......I suggest the authors incorporate additional representative baselines (at least one or two) during the rebuttal period.*\n\nWe appreciate reviewer's suggestion on making the experiment section more comprehensive. The primary reason for choosing FNO and Dil-ResNet as main baselines is that FNO is a represenative neural-PDE sovler that has shown decent performance on a wide array of problems, while\nDilResNet has shown good accuracy [1, 2] on the class of fluid problems investigated in this work. Nevertheless, following the reviewer's suggestion, we added experiments on Multi-wavelet NO, Factorized-FNO, and\nTensorized-FNO in our newly uploaded PDF (**Table 2,3,4**). In addition, we also added experiment results (**Table 1**) of our model on datasets from FNO,\nwhich are used in many other neural-PDE solver literatures. Due to the one page limit we only listed FNO/Linear Transformer/Dil-ResNet for reference, but we will incorpate more results from relevant literatures in the future.\n\n[1] Learned Simulators for Turbulence, ICLR 2023\n\n[2] Towards Multi-spatiotemporal-scale Generalized PDE Modeling, 2023\n\n---\n> *The use of notations in the paper is a little confusing. Does 'd' represent the dimension of the hidden layer? How is it selected?\n> 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.*\n\nYes, 'd' can be considered as the hidden dimension in the network. We set it to 128 by default and we invite reviewer to check out **Figure 3** in our **Appendix** where we show that the model's performance\ncan be further improved if we adopt a larger dimension (*We apologize for a typo in Figure 3's x-axis, the correct labelling should be: 128 -> 64, 64 -> 128*). We also added an illustrative diagram in the newly uploaed PDF (**Figure 2**) with annotated tensor shapes.\n\n\n---\n> *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?*\n\nIt is true that the 'v' has a shape that grows exponentially and we actually acknowledged it in the Conclusion section that our model is still not free from the curse of dimensionality. Our main goal is to alleviate the curse of dimensionality when using attention to parametrize the learnable kernel integral. More specifically, we exploit the low-rank structure of full-attention matrix and replace it with a set of much smaller attention matrices, which improves the computational efficiency and also the numerical stability when applying to higher-dimensional problems with a large number of grid points. \n\n---\n> *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.*\n\nWe thank reviewer's suggestion to explore the theoretical property of the decomposition process. While SVD-based methods (e.g. Tucker/CP decomposition) come with a strong theoretical foundation, they are not used in our work as it is prohibitively expensive to compute the online SVD for high-dimensional input tensor in each layer during training iterations. Instead, the decomposition from high-dimensional tensor to vector is accomplished through the cheap (yet effective) learnable projector we proposed. Investigating the theoretical capacity of the proposed projection method can be an interesting future direction.\n\nFollowing this, we would like to also elaborate a bit more about the difference between existing works and our work in terms of how the kernel is parameterized and computed.\n\n* Tensorized-FNO: Apply low-rank decomposition (e.g. Tucker/CP/Tensor-train) to the dense weight in spectral convolution layer and store them in the factorized form.\n* Factorized-FNO: Apply 1D spectral convolution along each dimension separately.\n* FactFormer: Use learnable projection to project tensor into vectors along different axes, and then use these projected vectors to compute data-dependent kernel.\n\nIn general, our work explored a new way to compute *data-dependent* kernel in a multi-dimensional factorized way.\n\n---\n> *I suggest you to include the following closely related work in the reference if you have not mentioned them in the paper...*\n\nWe thank the reviewer for the valuable references and we will include them in the future version of the manuscript (and camera-ready version if the paper gets accepted)."}]

Author rebuttal:
We would like to thank all the reviewers for their efforts spent reviewing our work and the suggestions on helping us improve our work. 

Below we summarize the newly added results that are included in the one-page pdf. 

#### **Table 1**: 
We tested our proposed model and Dilated-ResNet on two benchmark problems from FNO [1], which are 2D Navier-Stokes equation with different viscosities and 2D Darcy flow with different resolutions.
We also included the result of Linear Transformer [2] for comparison. In general, we observe that our proposed FactFormer also has competitive performance on these problems.
#### **Table 2/3/4**:
We added three new baselines on problems studied in the paper. The newly added baselines are Multi-wavelet neural operator(MWT) [3], and two newly proposed variants of FNO:
Tensorized-FNO [4], Factorized-FNO [5]. Tensorized-FNO(T-FNO) factorizes the weight in a standard FNO layer with tensor decomposition technique. Factorized-FNO (F-FNO) introduces a group of tricks for improving standard FNO and
applying spectral convolution in an axial-factorized way.

*Note: We opt for not including the results of MWT and T-FNO on 3D problems to avoid misleading conclusions.
 The current MWT implementation does not include architecture for applying 3D wavelet transform/bases projection and can only operate on resolutions that are a power of 2.
 The original T-FNO in [4] is applied to 1D/2D steady-state prediction problems. Our direct application of T-FNO on 3D time-dependent problems has a notable degradation compared to FNO.*
#### **Table 5**:
Detailed runtime comparison between different baseline models.
#### **Figure 1**:
Averaged rollout error trend of model using  linear/factorized attention scheme.
#### **Figure 2**:
A diagram comparing Axial Transformer[6] and the proposed model.

We have also addressed the detailed comments of every reviewer in separate replies. Please don’t hesitate to let us know if there is any further question or additional comment.

*References*:

[1] Fourier neural operator for parametric partial differential equations, 2021. 

[2] Choose a transformer: Fourier or galerkin, 2021. 

[3] Multiwavelet-based operator learning for differential equations, 2021. 

[4] Multi-Grid Tensorized Fourier Neural Operator for High Resolution PDEs, 2023. 

[5] Factorized Fourier Neural Operators, 2022. 

[6] Axial Attention in Multidimensional Transformers, 2019.


Meta-review: After rebuttal and discussions, the reviewers overall lean towards acceptance as the authors have addressed most of the concerns. Given the importance of these discussions and additional experiments in the acceptance of the paper, please carefully finalize them and incorporate the new results, discussions and citations in the paper. The main concern that remains is over the experimental setting. Please carefully report the codebase that was use and the hyper-parameters in all cases.
Decision: Accept (poster)