Title: Scalable Transformer for PDE Surrogate Modeling
Venue: NeurIPS 2023
Accepted: yes
Keywords: Efficient attention

Neural PDE solver
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

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 × 256 grid and 3D smoke buoyancy on a 64 × 64 × 64 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.

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 a wide range of scientific and engineering 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 a discretized grid and reduce the continuous differential equations to algebraic equations via methods like finite difference/element/volume methods or spectral method. Despite the theoretical guarantees behind them, their practical realization of specific problems can pose challenges that require careful expertise to overcome, such as a sufficient understanding of the underlying physics, or a fine-tailored mesh that resolves the necessary spatio-temporal scales. The interest in developing user-friendly and efficient PDE solvers, along with the success of deep learning models in many other areas [13, 44, 51, 103], has facilitated the emergence of neural-network-based PDE solvers, where the neural network can be used to parameterize the solution function of the target equation [96], or to 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.

Among various neural network designs, attention-based models (Transformer) [115] have become state-of-the-art for a 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 as 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 sequence-to-sequence Transformer [115] can be modified correspondingly to be better suited for

Pred.

Ref.

Time

2.00

1.75

10

1.50

1.25

1.00

10

0.75

20

0.50

0.25

-30

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 to the number of dimensions, this results in a very large attention matrix that computes the interaction between every pair of the grid points (despite this attention matrix is not evaluated explicitly in linear attention). Consequently, cascading a deep stack of attention layers introduces instability and relatively high computational cost on high-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-free attention, with a factorized integration scheme motivated by the inherent low-rank structure of dot-product kernel matrix. More specifically, we propose a multi-dimensional factorized kernel integral with each kernel function in the integral having only single-dimensional domains. To calculate these axial kernels, we propose a learnable integral operator that is able to project the input function with high-dimensional 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, our proposed model can scale up to multi-dimensional problems with a large number of grid points and achieve competitive performance compared to state-of-the-art models. Moreover, we show that our factorized attention mechanism can reduce the computational cost and improve stability compared to softmax-free linear attention. 1

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, or a 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 or partially known, such as large-scale climate modeling [61, 83, 90, 97]. The second group of work, namely the Physics-Informed Neural Networks (PINNs)[15, 40, 42, 52, 77, 87, 96, 106, 127], treat neural networks as 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 inputtarget data and can be trained solely based on 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 a 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

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 function [7, 108]. The training of neural operators can also be combined with the principle of PINNs to yield a 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 physics simulation, Transformer has drawn increasing research interest as a surrogate model for simulation, with its modeling capability demonstrated both as a neural PDE solver [16, 32, 35, 41, 43, 48, 54, 67, 82, 86] and as a pure data-driven model in the absence of a known governing PDE [14, 20, 31, 83]. The dot-product attention can be considered as an approximation of an integral transform with a non-symmetric learnable kernel function [16, 17, 35, 54, 60, 122], which relates Transformer to other popular operator learning models such the FNO [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 i } N q i =1 , keys { k i } N k i =1 , and values { v i } N v i =1 (assuming N k = N v ), attention mechanism [2, 33, 79, 115] dynamically computes a weighted average of the values: z i = ∑ N v j =1 h ( q i , k j ) v j , where q i , k i , v i ∈ R 1 × d , 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 i , k j ) = exp ( q i k T j /τ ) / ∑ s exp( q i k T s /τ ) , and τ is usually chosen as τ = √ d . The queries/keys/values are usually obtained from inputs via learnable projection. In self-attention, all of them are computed from the same source as follow:

where u i ∈ R 1 × d in is the input vector and { W q , W k , W v } ∈ R d in × d 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 as a 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 it as a Peterov-Galerkin 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.

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 } ∈ R N × d , their row vectors: q i / k i / v i , correspond to the sampling of a set of functions { q l ( · ) , k l ( · ) , v l ( · ) } d l =1 on grid point x i , where { x i } N i =1 discretizes the underlying domain. As a more concrete example, the l -th column (channel) of q i , represents the sampling of function q l ( · ) on a grid point, i.e. ( q i ) l = q l ( x i ) . Furthermore, softmax-free attention is equivalent to the numerical quadrature of a kernel integral:

where z i is the output vector, κ ( x, ξ ) = ∑ d l =1 q l ( x ) k l ( ξ ) is an instance-based kernel and w s is the quadrature weight. Understanding attention from the perspective of the kernel has been an active topic of research [17, 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 ( x i , ξ ) . 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 i , k j and their corresponding spatial coordinates x i , x j , RoPE [105] ( g ( · , · ) : R 1 × d × R ↦→ R 1 × d ) is defined as:

and λ, θ l are hyperparameters. θ l is usually chosen as 10000 -2( l -1) /d , l ∈ { 1 , 2 , . . . , d/ 2 } following Vaswani et al. [115] and Su et al. [105]. λ is a mesh-based weight that we heuristically set to 64 throughout most problems. The projection function Θ( · ) : R ↦→ R d × R d can explicitly modulate the dot product with relative position: g ( q i , x i ) g ( k j , x j ) T = q i Θ( x i -x j ) k T j , thanks to the following property of rotation matrix: R l ( x i ) R l ( x j ) T = R l ( x i -x j ) .

To summarize, we will adopt attention mechanism in the following form for the proposed model (with modification discussed in the next subsection):

where ˜ @ denotes a matrix whose row vectors are RoPE encoded as in (3), e.g., ˜ Q i = g ( q i , x i ) , w is the quadrature weight with a 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 ˜ Q ˜ K T evaluates the kernel function κ ( · , · ) on the discretization grid { x i } N i =1 .

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 a linear complexity by making use of the associativity of matrix multiplication (calculate ˜ K T V 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 1 × S 2 × . . . × S n = N points, the softmax-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 it is the product of two tall and thin matrices. Meanwhile, attending a large number of grid points to each other can be unstable and the linear attention has a complexity that is 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 aforementioned 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 a uniform Eulerian grid and can be treated as n -way tensor

Learnable

projection

V

Linear

Q/K

Linear

Attention

Q/K

Linear

Attention

Attention

Q/K

Linear

MLP

Input function

Output function

U ∈ R S 1 × S 2 × ... × S n 2 . The product of it with a matrix W ∈ R J × Sm across the m -th mode will result in a tensor of shape S 1 × . . . × S m -1 × J × S m +1 × . . . × S n , whose elements are defined as:

Learnable projection The first major component of the proposed framework is a set of learnable integral operators {G (1) , G (2) , . . . , G ( n ) } that projects the input function u : R n ↦→ R d into a set of functions with one-dimensional domain { ϕ (1) , ϕ (2) , . . . , ϕ ( n ) } ∈ R ↦→ R d , which is defined as:

where h ( m ) ( · ) : R d ↦→ R d and γ ( m ) ( · ) : R d ↦→ R d are pointwise learnable functions and w = 1 / ( L 1 L 2 · · · L m -1 L m +1 · · · L n ) , with L m being the size of domain Ω m discretized by { x ( m ) i } S m i =1 . In practice, we implement h ( m ) as a three-layer multi-layer perception (MLP) similar to the feedforward network in Transformer [115] and γ ( m ) 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 h ( 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):

where kernels { κ (1) , . . . , κ ( n ) } : R × R ↦→ R are computed based on projected single-dimensional function along each axis, v ( · ) : R n ↦→ R d is derived from the input function u via linear transformation (just as the value in standard attention). Next, we will discuss how the above kernel integral is implemented in practice. Using the learnable projection operator defined in (6), we can obtain { ˆ U (1) , . . . , ˆ U ( n ) } ( ˆ U ( m ) ∈ R Sm × d ) from the input U ∈ R S 1 × ... × Sn × d , where the i m -th row of ˆ U is the evaluation of projected function ϕ ( m ) ( · ) at x i m : ˆ U ( m ) i m = ϕ ( m ) ( x i m ) . Then we apply linear transformation on them to obtain the query/key matrix just as standard attention: Q ( m ) = ˆ U ( m ) W ( m ) q , K ( m ) = ˆ U ( m ) W ( m ) k , where { W ( m ) q , W ( m ) k } ∈ R d × d are learnable matrices. The query and key are used to compute the kernel matrix A ( m ) ∈ R S m × S m :

where w m is the mesh weight, ˜ @ denotes the RoPE encoded matrix as discussed in (4), the i -th row and j -th column of A ( m ) represents the kernel value κ ( m ) ( x ( m ) i ,x ( m ) j ) . Despite (8) has a quadratic complexity with respect to the grid size S m , this is an affordable cost for most of the problems where the axial grid size S m is mostly between 64 to 512 . Meanwhile, the value V ∈ R S 1 × ... × Sn × d is derived from the input via a linear transformation (i.e. ( m +1) -th mode product): V = U × n +1 W v , where W v ∈ R d × d is again a learnable matrix. The overall factorized kernel integral is numerically approximated with the following tensor-matrix product (Figure 2):

In (9), the computation of all kernel is of complexity O ( S 2 1 d + S 2 2 d + . . . + S 2 n d ) , and the time complexity of a single tensor-matrix product V × m A ( m ) is O ( NS m d ) . After evaluating the tensormatrix product, the output tensor Z will be sent to a pointwise feedforward network f ( · ) : R d ↦→ R d . To sum up, the update protocol of a single layer in our proposed factorized Transformer is defined as follows:

where Att ( · ) is the attention from (9), IN ( · ) is instance normalization [112] that normalizes each channel instance-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 O ( NS m d ) complexity (recall N = S 1 × . . . × S n ) instead of O ( S 2 m d ) as in our model. And its overall computation of attention is 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 row-wise manner and then the second block will conduct attention in a column-wise manner, whereas our model decomposes attention along all axes into a tensor-matrix product within every layer. We provide an illustrative example in Figure 27 of the Appendix.

3.3 Training techniques

In this subsection, we will discuss several techniques used for training the model (including 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 a simple pointwise learnable function ε ( · , · ) ∈ R d × R > 0 ↦→ R d can be used to propagate dynamics in the latent space with a fixed time interval ∆ t : z ( x, t +∆ t ) = z ( x, t ) + ε ( z ( x, t ) , t ) , where z is the output of the final attention layer. In practice, ε is implemented as a pointwise MLP and is efficient to compute. Leveraging this technique, with one call to the neural solver, we can forward the state for multiple time steps (by marching in the latent space for k steps), thus reducing the total number of calls by a ratio of k . This is in principle similar to the Temporal Bundling technique proposed in Message-Passing 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 { d 1 , d 2 , . . . , d k } and then adding them to the input u 0 by a forward Euler scheme in the physical space: ˆ u k = u 0 + d k ∆ t . 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.

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 in an unbounded rollout error. While there is hardly a universal method for guaranteeing their stability, a 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-walk 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, we compare our model against softmax-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 a 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 coarsegrained dynamics of turbulent flow and has demonstrated state-of-the-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 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

Wefirst 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 4 for smoke buoyancy similar to [37]. We also consider a well-known steady-state problem-2D Darcy flow, which has been studied in many of the previous works. Below we provide a 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 ω of the flow field within an interval [ t 0 , t 0 + T ] , where T = 1 s and t 0 is a random starting point in the sequence. We use a spatial grid of 256 × 256 and temporal discretization of ∆ t = 0 . 0625 s (therefore 1 s corresponds to 16 frames) to train and evaluate the model.

3D isotropic turbulence The second example is 3D isotropic turbulence governed by incompressible Navier-Stokes equation with a periodic boundary condition. The major difference from the first example is 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 Re λ = 84 [62]. The objective is to predict the pressure p and velocity u from t = 0 . 5 to t = 1 s (10 frames). The model is trained and evaluated on a 60 × 60 × 60 spatial grid with ∆ t = 0 . 05 s.

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 phiflow [46] 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 d and velocity of flow u from t = 3 to t = 15 s (16 frames). The model is trained and evaluated on a 64 × 64 × 64 spatial grid with ∆ t = 0 . 75 s. To account for non-periodic boundary conditions, we pad the domain for FNO variants and DilResNet following the original works. For FactFormer, we append a simple CNN block after the model, which comprises 3 -by3 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.

4.2 Results and discussion

For all the models, we study two protocols of training. The first is Latent Marching with Pushforward (denote as 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 k future steps. On 2D Kolmogorov flow/3D smoke buoyancy, k is set to 4 , and 2 for 3D isotropic turbulence. We interleave pushforward training with standard per-step training for LM models. The relative L 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 a 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 a 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 is 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 as 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 is highly dependent 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 a major difference between CNN-based models and neural operators.

-2)

4.50

FactFormer

10

*

FNO

×)

F-FNO

2.50

GalerkinTransformer

norm

*

DilResNet

1.50

2

Relative

1.00

0.75

0.50

141

211

421

Resolution

0.6

0.5

Linear

norm

Factorized

0.4

Relative

0.3

0.2

0.1

0.0

0.2

0.4

0.6

0.8

1.0

Prediction time (s)

4.3 Comparison against full attention

In this subsection, we will present an 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, we 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 · ,j || 2 = 1 . To accommodate for the memory cost of linear attention, we further downsample the 2D Kolmogorov flow discussed in the last subsection to a 128 × 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 of the two attention mechanisms in Figure 4 and Table 4. While in principle full attention could have better approximation capacity than factorized attention, in practice we find that it performs worse than factorized attention on this problem we studied. Specifically, its rollout is less stable and results in a degraded accuracy. We hypothesize that this is due to the instability of iteratively calculating the attention matrix of a large size, as 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 T is going to have a low-rank structure since rank ( A ) ≤ min( rank ( Q ) , rank ( K )) , Q,K ∈ R N × d and usually N >> 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 as the sum of singular values E = ∑ i σ i , where σ i is the i -th singular value and report the normalized cumulative energy histogram b k = ∑ k i =1 σ i / ∑ i σ i in Figure 5a, 5b, 5c. 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% of 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, 120]. Note that the spectrum of linear attention is based on a truncated SVD and therefore its rank will be even lower if a full SVD is performed. The highly low-rank structure of 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.

(a) Full attention matrix

A

(b) Attention matrix

A

(1)

(c) Attention matrix

A

(2)

1.0

1.0

1.0

0.9

energy

0.8

energy

0.8

Cumulativeenergy

0.8

0.7

Cumulative

0.6

Cumulative

0.6

0.6

0.4

0.5

0.4

0.2

=

First layer

Last layer

0.4

0.3

Lastlayer

First layer

0.2

Lastlayer

First layer

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Fraction of total singularvalues

Fractionoftotalsingularvalues

Fractionof totalsingularvalues

5 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 is still not free from the curse of dimensionality. The computation of the factorized kernel integral requires evaluating the function on all S 1 × S 2 × . . . × S m grid points. A future direction could be extending the factorization scheme to a more efficient tensor decomposition format like tensor-train. The proposed model currently exploits the uniform structure of the underlying grids and use mean pooling when doing projection, but non-uniform quadrature weight will be necessary when applying to 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.

References

Large-Scale Dynamics and Transition to Turbulence in the Two-Dimensional Kolmogorov Flow , pages 374-396. doi: 10.2514/5.9781600865831.0374.0396. URL https://arc.aiaa. org/doi/abs/10.2514/5.9781600865831.0374.0396 .

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2014. URL https://arxiv.org/abs/1409.0473 .

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-15349, 2019.

Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020.

Kaushik Bhattacharya, Bamdad Hosseini, Nikola B. Kovachki, and Andrew M. Stuart. Model reduction and neural networks for parametric pdes, 2021.

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.

Nicolas Boullé, Seick Kim, Tianyi Shi, and Alex Townsend. Learning green's functions associated with time-dependent partial differential equations. Journal of Machine Learning Research , 23(218):1-34, 2022.

Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for pde modeling. arXiv preprint arXiv:2209.04934 , 2022.

Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural PDE solvers. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 2241-2256. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/ v162/brandstetter22a.html .

Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint arXiv:2202.03376 , 2022.

Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K. Gupta. Clifford neural layers for pde modeling, 2023.

Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers, 2023.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:18771901, 2020.

Salva Rühling Cachay, Peetak Mitra, Haruki Hirasawa, Sookyung Kim, Subhashis Hazarika, Dipti Hingmire, Phil Rasch, Hansi Singh, and Kalai Ramea. Climformer-a spherical transformer model for long-term climate projections.

Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis. Physics-informed neural networks (pinns) for fluid mechanics: A review. Acta Mechanica Sinica , 37(12):1727-1738, 2021.

Shuhao Cao. Choose a transformer: Fourier or galerkin. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 24924-24940. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ d0921d442ee91b896ad95059d13df618-Paper.pdf .

Shuhao Cao, Peng Xu, and David A. Clifton. How to understand masked autoencoders, 2022.

Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers, 2020.

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.1017/jfm.2013.122.

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.

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.1109/72.392253.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers, 2019.

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.

Mengyu Chu and Nils Thuerey. Data-driven synthesis 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://doi.org/10.1145%2F3072959.3073643 .

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth, 2021.

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.11929 , 2020.

Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zepeda-Núñez, Jamie A. Smith, Michael P. Brenner, and Stephan Hoyer. Learning to correct spectral methods for simulating turbulent flows, 2022.

Antonio H de O Fonseca, Emanuele Zappala, Josue Ortega Caro, and David van Dijk. Continuous spatiotemporal transformers. arXiv preprint arXiv:2301.13338 , 2023.

Kai Fukami, Koji Fukagata, and Kunihiko Taira. Super-resolution reconstruction of turbulent flows with machine learning. Journal of Fluid Mechanics , 870:106-120, 2019. doi: 10.1017/ jfm.2019.238.

Zhihan Gao, Xingjian Shi, Hao Wang, Yi 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-25403, 2022.

Nicholas Geneva and Nicholas Zabaras. Transformers for modeling physical systems. Neural Networks , 146:272-289, 2022.

Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines, 2014. URL https: //arxiv.org/abs/1410.5401 .

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.13587 , 2021.

Ruchi Guo, Shuhao Cao, and Long Chen. Transformer meets boundary value inverse problems. arXiv preprint arXiv:2209.14977 , 2022.

Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differential equations, 2021.

Jayesh K. Gupta and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling, 2022.

Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv preprint arXiv:2209.15616 , 2022.

Nathan Halko, Per-Gunnar Martinsson, and Joel A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions, 2010.

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.

Xu Han, Han Gao, Tobias Pfaff, Jian-Xun Wang, and Li-Ping Liu. Predicting physics in mesh-reduced space with temporal attention. arXiv preprint arXiv:2201.09113 , 2022.

Zhongkai Hao, Songming Liu, Yichi Zhang, Chengyang Ying, Yao Feng, Hang Su, and Jun Zhu. Physics-informed machine learning: A survey on problems, methods and applications, 2023.

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.

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.

Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers, 2020. URL https://openreview.net/forum?id= H1e5GJBtDr .

Philipp Holl, Vladlen Koltun, and Kiwon Um. phiflow: A differentiable pde solving framework for deep learning via physical simulations.

Ameya D Jagtap and George Em Karniadakis. Extended physics-informed neural networks (xpinns): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations. Communications in Computational Physics , 28(5): 2002-2041, 2020.

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 transformers. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=mfIX4QpsARJ .

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 framework. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-15, 2020. doi: 10.1109/SC41405.2020.00013.

Pengzhan Jin, Shuai Meng, and Lu Lu. Mionet: Learning multiple-input operators via tensor product. SIAM Journal on Scientific Computing , 44(6):A3490-A3514, 2022.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics , 3(6):422-440, 2021.

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://proceedings.mlr.press/v119/katharopoulos20a.html .

Georgios Kissas, Jacob Seidman, Leonardo Ferreira Guilhoto, Victor M. Preciado, George J. Pappas, and Paris Perdikaris. Learning operators with coupled attention, 2022.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations , 2020. URL https://openreview. net/forum?id=rkgNKkHtvB .

Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine learning-accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences , 118(21), may 2021. doi: 10.1073/pnas.2101784118. URL https://doi.org/10.1073%2Fpnas.2101784118 .

Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications. SIAM Rev. , 51(3):455-500, aug 2009. ISSN 0036-1445. doi: 10.1137/07070111X. URL https: //doi.org/10.1137/07070111X .

Jean Kossaifi, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high resolution PDEs, 2023. URL https: //openreview.net/forum?id=po-oqRst4Xm .

Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators, 2021.

Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint arXiv:2108.08481 , 2021.

Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-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.

A. 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.1063/1.1833415. URL https://doi.org/10.1063/1.1833415 . 015106.

Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition, 2015.

Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids, 2019.

Zijie Li and Amir Barati Farimani. Graph neural network-accelerated lagrangian fluid simulation. Computers & Graphics , 103:201-211, 2022. ISSN 0097-8493. doi: https: //doi.org/10.1016/j.cag.2022.02.004. URL https://www.sciencedirect.com/science/ article/pii/S0097849322000206 .

Zijie Li, Tianqin Li, and Amir Barati Farimani. TPU-GAN: Learning temporal coherence from dynamic point cloud sequences. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=FEBFJ98FKx .

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://openreview.net/forum?id=EPPqt3uERT .

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.08895 , 2020.

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.

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.

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.03485 , 2020.

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.

Xinliang Liu, Bo Xu, and Lei Zhang. Ht-net: Hierarchical transformer based operator learning model for multiscale pdes. arXiv preprint arXiv:2210.10890 , 2022.

Xinliang Liu, Bo Xu, and Lei Zhang. Mitigating spectral bias for the multiscale operator learning with hierarchical attention, 2023.

Winfried Lötzsch, Simon Ohler, and Johannes Otterbach. Learning the solution operator of boundary value problems using graph neural networks. In ICML 2022 2nd AI for Science Workshop , 2022. URL https://openreview.net/forum?id=4vx9FQA7wiC .

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.03193 , 2019.

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.

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: 10.1016/j.cma.2022.114778. URL https: //doi.org/10.1016%2Fj.cma.2022.114778 .

Minh-Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation, 2015. URL https://arxiv.org/abs/1508. 04025 .

Xindian Ma, Peng Zhang, Shuai Zhang, Nan Duan, Yuexian Hou, Dawei Song, and Ming Zhou. A tensorized transformer for language modeling, 2019.

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.1016/j.cpc.2016.02.005. URL https://doi.org/10.1016%2Fj.cpc.2016.02.005 .

Tan Nguyen, Minh Pham, Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. Fourierformer: Transformer meets generalized fourier integral theorem. Advances in Neural Information Processing Systems , 35:29319-29335, 2022.

Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. Climax: A foundation model for weather and climate, 2023.

Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov. Tensorizing neural networks, 2015.

I. V. Oseledets. Tensor-train decomposition. SIAM Journal on Scientific Computing , 33(5):2295-2317, 2011. doi: 10.1137/090752286. URL https://doi.org/10.1137/ 090752286 .

Oded Ovadia, Adar Kahana, Panos Stinis, Eli Turkel, and George Em Karniadakis. Vito: Vision transformer-operator. arXiv preprint arXiv:2303.08891 , 2023.

Guofei Pang, Lu Lu, and George Em Karniadakis. fpinns: Fractional physics-informed neural networks. SIAM Journal on Scientific Computing , 41(4):A2603-A2626, 2019.

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.1063/5.0062546. URL https://doi.org/10.1063%2F5. 0062546 .

Jaideep Pathak, Mustafa Mustafa, Karthik Kashinath, Emmanuel Motheau, Thorsten Kurth, and Marcus Day. Using machine learning to augment coarse-grid computational fluid dynamics simulations. arXiv preprint arXiv:2010.00072 , 2020.

Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, 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.

Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W. Battaglia. Learning mesh-based simulation with graph networks, 2021.

Lukas Prantl, Benjamin Ummenhofer, Vladlen Koltun, and Nils Thuerey. Guaranteed conservation of momentum for learning particle-based fluid dynamics, 2022.

Alfio Quarteroni and Alberto Valli. Domain decomposition methods for partial differential equations . Number BOOK. Oxford University Press, 1999.

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 Systems , volume 20. Curran Associates, Inc., 2007. URL https://proceedings.neurips.cc/ paper_files/paper/2007/file/013a006f03dbc5392effeb8f18fda755-Paper.pdf .

Md Ashiqur Rahman, Zachary E. Ross, and Kamyar Azizzadenesheli. U-no: U-shaped neural operators, 2023.

Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics , 378:686-707, 2019.

Stephan Rasp, Peter D. Dueben, Sebastian Scher, Jonathan A. Weyn, Soukayna Mouatadid, and Nils Thuerey. WeatherBench: A benchmark data set for data-driven weather forecasting. Journal of Advances in Modeling Earth Systems , 12(11), nov 2020. doi: 10.1029/2020ms002203. URL https://doi.org/10.1029%2F2020ms002203 .

Robert S. Rogallo. Numerical experiments in homogeneous turbulence. 1981.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models, 2022.

Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks, 2020.

Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities, 2020.

Dule Shu, Zijie Li, and Amir Barati Farimani. A physics-informed diffusion model for high-fidelity flow field reconstruction. Journal of Computational Physics , 478:111972, 2023.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587): 484-489, 2016.

Kimberly Stachenfeld, Drummond B. Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez-Gonzalez. Learned coarse models for efficient turbulence simulation, 2022.

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2022.

Luning Sun, Han Gao, Shaowu Pan, and Jian-Xun Wang. Surrogate modeling for fluid flows based on physics-constrained deep learning without simulation data. Computer Methods in Applied Mechanics and Engineering , 361:112732, 2020.

Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains, 2020.

Jingwei Tang, Vinicius C Azevedo, Guillaume Cordonnier, and Barbara Solenthaler. Neural green's function for laplacian systems. Computers & Graphics , 107:186-196, 2022.

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.

Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators, 2023.

Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: A unified understanding of transformer's attention via the lens of kernel, 2019.

Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization, 2017.

Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey. Solver-in-theloop: Learning from differentiable physics to interact with iterative pde-solvers. Advances in Neural Information Processing Systems , 33:6111-6122, 2020.

Benjamin Ummenhofer, Lukas Prantl, Nils Thuerey, and Vladlen Koltun. Lagrangian fluid simulation with continuous convolutions. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=B1lDoJSYDH .

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. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper/2017/file/ 3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physicsinformed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , KDD '20, page 1457-1466, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 9781450379984. doi: 10.1145/3394486.3403198. URL https://doi.org/10.1145/3394486.3403198 .

Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed deeponets. Science Advances , 7(40): eabi8605, 2021. doi: 10.1126/sciadv.abi8605. URL https://www.science.org/doi/ abs/10.1126/sciadv.abi8605 .

Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed deeponets, 2021.

Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks, 2022.

Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

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:104180, 2022.

Matthew A Wright and Joseph E Gonzalez. Transformers are deep infinite-dimensional non-mercer binary kernel machines. arXiv preprint arXiv:2106.01506 , 2021.

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.

Yinchong Yang, Denis Krompass, and Volker Tresp. Tensor-train recurrent neural networks for video classification, 2017.

Fisher Yu and Vladlen Koltun. Multi-scale context aggregation by dilated convolutions, 2016.

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.

Min Zhu, Handi Zhang, Anran Jiao, George Em Karniadakis, and Lu Lu. Reliable extrapolation of deep neural operators informed by physics or sparse observations. Computer Methods in Applied Mechanics and Engineering , 412:116064, 2023.

A Model implementation details

The major hyperparameters are listed in Table 5.

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 i ← z i + ψ ( x i ) , where x i is the Cartesian coordinate of latent encoding z i , ψ : R n ↦→ R d 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 d k , the number of function bases used to compute the kernel: κ ( x, ξ ) = ∑ d k l q l ( x ) k l ( ξ ) . We train the model with AdamW optimizer and cyclic learning rate scheduler with a maximum learning rate 3 e -4 , similar to prior Transformer-for-PDE 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 a U-shape arrangement with 4 CNN layers, with all layers using a kernel size of 3 and padding size of 1 (pad with zeros). The first CNN layer has a stride of 2, while other layers have a 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 a 2D convolutional layer to compress the temporal dimension if it is a time-dependent problem. Concretely, we first reshape the input into ( N,T in ) where N is the number of spatial grid points and T in is the number of input frames. Then we apply 2D convolution filters of size (1 , T in ) to compress the temporal dimension to 1 . At the bottom of every model, we use a three-layer MLP 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 implementation: https://github.com/neuraloperator/physics_informed . And for Factorized Fourier Neural Operator the implementation is taken from https://github.com/alasdairtran/ fourierflow . We add group normalization before the final fully-connected layer. We use a 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.

For Dil-ResNet, we adopt the implementation from Gupta and Brandstetter [37]: https://github. com/microsoft/pdearena . Compared to Gupta and Brandstetter [37] and Stachenfeld et al. [104], we simplify the setting by truncating the number of layers inside each block, where we use dilation layers [1 , 3 , 8 , 3 , 1] ( [1 , 2 , 4 , 2 , 1] is chosen for the 3D problem as it has slightly better performance) inside each block instead of [1 , 2 , 4 , 8 , 4 , 2 , 1] . The primary reason is that without truncation, the training on 3D problems will take over a week (and cannot fit into a single A6000 GPU for 2-step

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-based (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://github.com/BaratiLab/OFormer , with a 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 5 e -4 to 5 e -6 throughout the training.

All the experiments are carried out using PyTorch 1.8 except for FNO/F-FNO experiment, which uses PyTorch 1.13 for optimizing complex-valued 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.

0.6

DilRes

FNO

0.5

FactFormer

F-FNO

norm

0.4

J

Relative

0.3

R

0.2

0.1

0.0

0.2

0.4

0.6

0.8

1.0

Prediction time (s)

0.6

DilRes

FNO

0.5

FactFormer

F-FNO

Relative L2 norm

0.4

0.3

0.2

0.1

0.0

0.1

0.2

0.3

0.4

0.5

Prediction time (s)

0.6

DilRes

FNO

0.5

FactFormer

F-FNO

norm

0.4

Relative

0.3

R

0.2

0.1

0.0

0.1

0.2

0.3

0.4

0.5

Prediction time(s)

0.6

DilRes

FNO

0.5

FactFormer

F-FNO

norm

0.4

5

Relative

0.3

R

0.2

0.1

0.0

2

4

6

8

10

12

Prediction time (s)

D Further ablation study

This section includes further study and numerical experiments on the proposed model.

Note on backpropagation Consider a single (softmax-free) attention block that consists of a attention layer and a two layer feedforward network (for simplicity we consider single-headed case):

where U 0 ∈ R N × d is the input, W 1 , W 2 ∈ R d × d are learnable weights for the feedforward network. Given a loss function l ( · ) : R N × d ↦→ R (for example, mean squared error), define ˜ U 1 := U 1 W 1 , ∂l ∂U 2 := h , ∂σ ( ˜ U 1 ) ∂ ˜ U 1 := g , the gradient for the weights and input in the feedforward network are ( ⊙ denotes element-wise multiplication):

For linear (softmax-free) dot-product attention: U 1 = Att ( U 0 ) = QK T V , where Q = U 0 W q , K = U 0 W k , V = U 0 W v , and thus U 1 = U 0 W q W T k U T 0 U 0 W v , the gradient of weights are:

where U 0 ∈ R N × d has appeared three times in each calculation and thus in backpropagation the computational cost of attention layer is generally more expensive than the feedforward layer where it 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, it can be written as:

0.6

DilRes

FNO

0.5

FactFormer

F-FNO

norm

0.4

Relative

0.3

R

0.2

0.1

0.0

2

4

Prediction time (s)

6

8

10

12

where Z i,c denotes the i -th row and j -th column of matrix Z and similar for other matrices, Q = XW q , K = XW k , V = XW v and X ∈ R N × d is input. The gradient of parameters are computed as:

For the proposed factorized attention, it can written as:

where A ( m ) = Q ( n ) ( K ( m ) ) T ( A ( m ) ∈ R S m × S m , N = S 1 × S 2 × . . . × S n ) is the axial kernel matrix as defined in (8), X ( m ) ∈ R S m × d is the axial projection along m -th axis as defined in (6).

For W v , its gradient is computed as:

For W ( n ) q , W ( n ) k , their gradients are computed as:

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.

Runtime comparison As discussed in Section A, the kernel dimension indicates how many function bases are used to evaluate the kernel and a larger kernel dimension is beneficial to the learning capacity of the model. As shown in Figure 11a, 11b, linear attention's training cost increases more significantly than the factorized attention as kernel dimension increases, since its complexity is quadratic with respect to kernel dimension. Factorized attention's computational efficiency can be further improved by reducing the spatial resolution, leveraging techniques such as learning the mapping in the latent space (similar to latent diffusion model [99]), multi-scale network architecture that resembles multigrid methods [35, 74], or domain decomposition [47, 93]. Furthermore, the training cost of factorized attention is also relatively lower than linear attention on 3D domain as shown in Figure 12a, 12b.

Factorized

Linear

25000

OOM

20000

19209

(qw)

snow

15000

15620

12029

10000

8439

5161

5217

5275

5329

5553

5000

64

128

192

256

512

Kernel dimension

(a) Forward + backward peak memory usage.

Factorized

Linear

OOM

0.5

0.453

Fwd.+Bwd. Time (sec)

0.4

0.378

0.301

0.3

0.232

0.2

0.1

0.093

0.094

0.096

0.097

0.102

0.0

64

128

192

256

512

Kernel dimension

(b) Forward + backward time per iteration.

Factorized

Linear

50000

OOM

40000

38679

(qw)

34468

Memory usage

30000

20000

14593

11498

10000

4375

0

403

603

803

Grid size

(a) Forward + backward peak memory usage.

Factorized

Linear

OOM

2.00

1.75

1.Time (sec)

1.50

1.359

1.25

Fwd.+Bwd.

1.00

0.864

0.75

0.50

0.364

0.379

0.25

0.121

0.00

403

603

803

Grid size

(b) Forward + backward time per iteration.

Model scaling performance We study the impact of the number of heads and size of kernel dimension on 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 a 128 × 128 grid) with a splitting different from the Evaluation Section in the main body of the paper. As shown in Figure 13b, the number of attention heads has a crucial impact on the final performance. The model's performance drops significantly when using fewer heads. This highlights the importance of multi-head mechanism in the factorized attention. In addition, we observe that the final accuracy of our model benefits from an increased kernel dimension (as shown in Figure 13a).

(a) Ablation on the kernel dimension.

(b) Ablation on the number of heads.

0.183

0.154

loss

0.15

0.131

0.119

2

Relative

0.10

R

0.05

0.00

64

128

256

512

Kerneldimension

0.510

0.5

0.420

loss

0.4

2

0.3

Relative

0.209

0.2

0.154

0.1

0.0

1

2

4

8

Numberofheads

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 κ ( ξ 1 , ξ 2 ) depends only on the relative distance between two points, e.g. L 2 distance: || ξ 1 -ξ 2 || 2 ). The kernel matrices also exhibits symmetric pattern despite the non-symmetric nature of dot product QK T and all kernel matrices are diagonal dominated.

(a) Layer 1, Head 1:

kernel of

y

axis

(b) Layer 1, Head 1:

kernel of

x

axis

(c) Layer 4, Head 1:

kernel of

y

axis

(d) Layer 4, Head 1:

kernel of

x

axis

51

51

51

&1

0.00

0.00

1.57

3.14

4.71

6.28

1.00

0.00

0.00

1.57

3.14

4.71

6.28

1.00

0.00

00'0

1.57

3.14

4.71

6.28

1.00

0.00

0.00

1.57

3.14

4.71

6.28

1.00

0.75

- 0.75

0.75

0.75

1.57

0.50

1.57

0.50

1.57

0.50

1.57

05'0

0.25

0.25

3.14

0.00

E

0.00

3.14

0.00

E

0.00

0.25

0.25

0.25

0.25

4.71

0.50

4.71

0.50

4.71

0.50

4.71

0.50

0.75

0.75

0.75

0.75

6.28

-1.00

6.28

1.00

6.28

1.00

6.28

1.00

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 1

0.25

Seed2

Seed3

rm

0.20

nor

5

e

0.15

ative

R

0.10

0.05

0.2

0.4

0.6

0.8

1.0

Time (s)

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,

where ω denotes vorticity, u denotes velocity, Re denotes the Reynolds number, x = ( x 1 , x 2 ) denotes the spatial coordinates and f ( · ) is the forcing term that is set to f ( x ) = -n cos ( nx 2 ) -0 . 1 ω ( 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 8 [1, 19] and introduce dragging force term 0 . 1 ω ( x ) as described in Kochkov et al. [56]. The initial condition ω 0 is sampled from a 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://github.com/ neuraloperator/physics_informed/blob/master/solver/kolmogorov_flow.py to generate the data. The referenced direct numerical simulation is carried out with a spatial resolution of 2048 × 2048 and a temporal resolution of 1 e -4 .

3D isotropic turbulence The incompressible Navier-Stokes equation for this problem is given as:

where u denotes velocity, p denotes the pressure, ν is the viscosity parameter, x = [ x 1 , x 2 , x 3 ] 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 https://github.com/ spectralDNS . The referenced simulation is carried out with a spatial resolution of 60 × 60 × 60 and a temporal resolution of 0 . 005 s .

3D smoke buoyancy The governing equations for the 3D smoke buoyancy problem are incompressible Navier-Stokes equation (similar as above) coupled with advection equation:

where f ( x , t ) = [0 , 0 , ηd ( x , t )] , η is the buoyancy factor, the velocity field u has a Dirichlet boundary condition: u ( x , · ) = 0 , ∀ x ∈ ∂ Ω , and the scalar density field for smoke has a Neumann boundary condition: ∇ d ( x , · ) = 0 , ∀ x ∈ ∂ Ω . The initial condition d 0 ( x ) is a 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 is 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://github.com/microsoft/pdearena to generate the data. The solver applies an advection-project scheme. The referenced simulation is carried out with a spatial resolution of 64 × 64 × 64 and a temporal resolution of 0 . 75 s .

2D Darcy flow The equation for the 2D Darcy flow is defined as:

where f ( x ) is the forcing function that is set to 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 on a 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.

F 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=t0+0.2

t0+0.4

Reference

Error

t0+0.6

t0+0.8

t0+1.0

Prediction

20

12

10

T=t0+0.2

t0+0.4

Error

t0+0.6

t0+0.8

t0+1.0

Reference

Prediction

20

10

10

20

20.0

17.5

15.0

12.5

10.0

7.5

5.0

2.5

T=0.6

0.7

Error

0.8

0.9

1.0

Reference

Prediction

4.0

3.5

3.0

2.5

2.0

1.5

1.0

0.5

T=0.6

0.7

Error

0.8

0.9

1.0

Reference

Prediction

1.0

0.8

0.6

0.4

0.2

T=0.6

0.7

Error

0.8

0.9

1.0

Reference

Prediction

1.0

80-

0.6

0.4

0.2

T=0.6

0.7

Error

0.8

0.9

1.0

Reference

Prediction

1.2

1.0

0.8

0.6

0.4

0.2

T=6

8.25

Error

10.5

12.75

15

Reference

Prediction

2.00

1.75

1.50

1.25

1.00

0.75

0.50

0.25

0.40

0.35

0.30

0.25

0.20

0.15

0.10

0.05

T=6

8.25

Error

10.5

12.75

15

Reference

Prediction

0.3

0.2

0.1

0.0

0.1

0.2

0.3

0.35

0.30

0.25

0.20

0.15

0.10

0.05

T=6

8.25

Error

10.5

12.75

15

Reference

Prediction

0.3

0.2

0.1

0.0

0.1

0.2

0.3

0.25

0.20

0.15

0.10

0.05

0.00

T=6

8.25

Error

10.5

12.75

15

Reference

Prediction

0.4

0.3

0.2

0.1

0.0

0.1

0.2

0.3

0.4

0.4

0.3

0.2

0.1

0°0

Reference

Prediction

Error

0.0150

0.0125

0.0100

0.0075

0.0050

0.0025

0.0003

0.0002

0.0001

G Broader impact

The numerical simulation of PDEs is of extensive application in various fields, such as manufacturing, 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-based PDE surrogate models. Our experiments demonstrate that the proposed model serves as an efficient surrogate for numerical solvers, maintaining a balance between accuracy and efficiency, and thus pushing the Pareto front of accuracy-efficiency. However, as there exists a large variety of PDEs and each with very unique properties, there is no guarantee that one 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 to acknowledge the limitations and potential risks associated with the application of neural PDE solvers.

Apart from enriching the existing architecture design choice of attention-based models, our work also has the potential to 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:x

input: x

(size: [BS,h,w,d])

(size:[BS,h,w,d])

legend

transpose to

(1)

(2)

batch axis

BS

batch size

h

height

Φ(1)(x)

Φ(2)(x)

(size: [BS x h,w,d])

W

width

(size: [BS,h,d])

(size: [BS,w,d])

↑

↑

d

channel

emb

emb

emb

emb

Φ(k)

operator along

projection

thekth axis

emb point-wise linear

Q(1)

K(1)

Q(2)

K(2)

V

Q

K

Q,K,V size:

embedding

[BSxh,w,d]]

dot product

V

A(1) := Q(1)(K(1))

A(2) := Q(2)(k(2))7

(size:[BS × h,w,w])

(size: [BS,h,w,d])

(size: [BS,h,h])

(size:[BS,w,w])

product

+

V × A(1) ×2 A(2)

(QKT)V

K:andano

output:y

(size: [BS,h, w,d])

(size: [BS × h,w,d])

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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)\n\nQUESTIONS:\nNone.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-15"}

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

STRENGTHS:
NA

WEAKNESSES:
The standard Transformer has two components: attention and MLP. This paper only addresses the attention. But the MLP complexity O(Nd^2) would still make the overall complexity exponential in spatial dimension. This seems to defeat the purpose of the paper.

By design, the attention without softmax has rank d, where d is the hidden size and is typically much smaller than the sequence length N. This is shown in Fig 8a. For A^(1) and A^(2) in Fig 8b and 8c, the sequence length along each separate axis is much smaller. This gives the impression that matrices A^(1) and A^(2) are less rank deficient. But this does not mean that the proposed attention is better than vanilla attention.

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

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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)$.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe author properly acknowledged the curse of dimensionality. Moreover, the dimension-bound factorization limits the domain type to tensor-product type (i.e., rectangular).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

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

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

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

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

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nListed in weaknesses and questions sections.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1) 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.\n\nWEAKNESSES:\n1) (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\nQUESTIONS:\n1) 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\".\n\nLIMITATIONS:\nna", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-27"}