SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
Z = \operatorname {A t t} (U) = V \times_ {1} A ^ {(1)} \times_ {2} A ^ {(2)} \times \dots \times_ {n} A ^ {(n)}. \tag {9}
$$
In (9), the computation of all kernel is of complexity  $O(S_1^2 d + S_2^2 d + \ldots + S_n^2 d)$ , and the time complexity of a single tensor-matrix product  $V \times_m A^{(m)}$  is  $O(NS_m d)$ . After evaluating the tensor-matrix product, the output tensor  $Z$  will be sent to a pointwise feedforward network  $f(\cdot): \mathbb{R}^d \mapsto \mathbb{R}^d$ . To sum up, the update protocol of a single layer in our proposed factorized Transformer is defined as follows:
$$
U \leftarrow f (\operatorname {I N} (\operatorname {A t t} (U))) + U, \tag {10}
$$
where  $\mathrm{Att}(\cdot)$  is the attention from (9),  $\mathrm{IN}(\cdot)$  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} \times \ldots \times S_{n}$ ) instead of  $O(S_m^2 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  $\varepsilon (\cdot ,\cdot)\in \mathbb{R}^d\times \mathbb{R}_{>0}\mapsto \mathbb{R}^d$  can be used to propagate dynamics in the latent space with a fixed time interval  $\Delta t$ :  $z(x,t + \Delta t) = z(x,t) + \varepsilon (z(x,t),t)$ , where  $z$  is the output of the final attention layer. In practice,  $\varepsilon$  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,\dots ,d_k\}$  and then adding them to the input  $u_{0}$  by a forward Euler scheme in the physical space:  $\hat{u}_k = u_0 + d_k\Delta 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.
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 coarse-grained 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
We first apply our model to three fluid-like systems, where the underlying physics patterns are sensitive to the spatiotemporal scale that discretization can resolve, and typically require fine discretization for classical numerical solvers. In these problems, the neural PDE solver is trained to predict the next frame (or multiple frames if using latent marching) given a context of previous frames. The number of context frames of Kolmogorov flow and isotropic turbulence is set to 10 following [68, 72], and 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  $\omega$  of the flow field within an interval  $[t_0, t_0 + T]$ , where  $T = 1\mathrm{s}$  and  $t_0$  is a random starting point in the sequence. We use a spatial grid of  $256 \times 256$  and temporal discretization of  $\Delta t = 0.0625\mathrm{s}$  (therefore 1s 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_{\lambda} = 84$  [62]. The objective is to predict the pressure  $p$  and velocity  $\mathbf{u}$  from  $t = 0.5$  to  $t = 1$ s (10 frames). The model is trained and evaluated on a  $60 \times 60 \times 60$  spatial grid with  $\Delta 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  $\nu = 0.003$ . The objective is to predict the scalar density field of smoke  $d$  and velocity of flow  $\mathbf{u}$  from  $t = 3$  to  $t = 15s$  (16 frames). The model is trained and evaluated on a  $64\times 64\times 64$  spatial grid with  $\Delta t = 0.75s$ . 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-by-3 convolutional layers with zero padding.
2D Darcy flow In addition to the above time-dependent 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.
Model	FNO2D	F-FNO2D	Dil-ResNet	FactFormer
AR*	LM	AR*	LM	AR	LM	AR	LM
ω avg. error	0.3177	0.2978	0.1486	0.2453	0.8156	0.1655	0.8835	0.1734
ω final error	0.4423	0.4567	0.2811	0.3861	1.1692	0.3051	1.0963	0.3017
Inf. time (s)	0.73	0.81	0.86	1.01	4.69	1.78	3.14	1.38
# params (M)	85.1	3.7	2.4	3.5
Table 1: Evaluation results of 2D Kolmogorov flow. A batch size of 10 is used for inference. LM models predict 4 steps with each call to the model. Total prediction length is 16 steps. AR*: Since for 2D problem FNO variants can afford to rollout more steps during training, AR FNO rollout for 12 steps, AR F-FNO rollout for 6 steps, whereas other AR models rollout for 2 steps during training. For model that has complex parameters, each cffloat parameter count as two parameters.
Model	FNO3D	F-FNO3D	Dil-ResNet	FactFormer
AR	LM	AR	LM	AR	LM	AR	LM
p avg. error	0.8080	0.4634	0.3151	0.3264	0.1725	0.1778	0.2989	0.2545
p final error	1.1285	0.6522	0.4250	0.4159	0.2573	0.2448	0.4407	0.3431
u avg. error	0.3967	0.3382	0.2298	0.2303	0.1143	0.1250	0.1775	0.1670
u final error	0.6561	0.4735	0.2799	0.2850	0.1675	0.1671	0.2594	0.2218
Inf. time (s)	1.01	0.91	2.77	1.37	12.67	6.89	2.68	1.31
# params (M)	509.8	3.0	6.9	5.1
Table 2: Evaluation results of 3D isotropic turbulence. A batch size of 4 is used for inference. LM models predict 2 steps with each call to the model. Total prediction length is 10 steps.
Model	FNO3D	F-FNO3D	Dil-ResNet	FactFormer
AR	LM	AR	LM	AR	LM	AR	LM
d avg. error	0.1607	0.1344	0.1038	0.1236	0.0843	0.0999	0.1017	0.0942
d final error	0.1775	0.1287	0.1415	0.1219	0.1070	0.1062	0.1693	0.0941
u avg. error	0.5198	0.4255	0.3419	0.3713	0.2378	0.2747	0.3537	0.2592
u final error	1.0245	0.6718	0.8655	0.6146	0.5372	0.5023	0.7881	0.4482
Inf. time (s)	3.19	1.47	6.35	2.75	27.49	6.94	5.61	2.62
# params (M)	509.8	3.0	6.9	4.6
Figure 4: Error trend of different Transformer models on 2D Kolmogorov flow with  $128 \times 128$  grid.
Figure 3: Error on 2D Darcy flow with different training resolutions. Galerkin Transformer's result is taken from the original paper [16].
Table 3: Evaluation results of 3D smoke buoyancy. A batch size of 4 is used for inference. LM models predict 4 steps with each call to the model. Total prediction length is 16 steps.
Model	Avg. rel. L2Norm	Fwd.	Fwd. + Bwd.
Enc. time (s)	Prop. time (s)	Time (s)	Mem. (MB)
Factorized attention	0.1529	0.0202	0.0954	5217
Linear attention (matmul)	0.5853	0.1370	0.0112	0.3013	12029
Linear attention (einsum)	0.1333	0.2938	12029
Table 4: Comparison between factorized and linear attention on their forward/backward computational cost, Mem. denotes the peak memory usage. The benchmark is carried out using PyTorch 1.8.2 on an RTX 3090, with a batch size of 4. Enc. time: the time spent on obtaining the latent encoding, primarily includes attention layers and feedforward layers after each attention layer; Prop. time: the time used to propagate dynamics in the latent space with a 3-layer MLP.
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  $\tilde{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\times 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  $\mathrm{rank}(A) \leq \min(\mathrm{rank}(Q), \mathrm{rank}(K))$ ,  $Q, K \in \mathbb{R}^{N \times 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 = \sum_{i} \sigma_{i}$ , where  $\sigma_{i}$  is the  $i$ -th singular value and report the normalized cumulative energy histogram  $b_{k} = \sum_{i=1}^{k} \sigma_{i} / \sum_{i} \sigma_{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)}$
Figure 5: Spectrum of different attention matrices. (a): Attention matrices from softmax-free attention. (b), (c): Attention matrices for axis  $x$  and  $y$  from FactFormer. The vertical line indicates the fraction of singular values that capture  $90\%$  of the total energy. Since the full attention matrix  $A$  has a relatively large size  $(16384 \times 16384)$ . Therefore its spectrum is computed via TruncatedSVD [39] with top 1024 components truncated.
(c) Attention matrix  $A^{(2)}$
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} \times S_{2} \times \ldots \times 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.
This work is supported by the National Science Foundation under Grant No. 1953222. The authors would like to thank the anonymous reviewers and area chair for their efforts and valuable feedback during the reviewing process. The authors would like to thank Dr. Shuhao Cao from University of Missouri - Kansas City for the comments regarding the backpropagation of feedforward layer and attention layer in Transformer. The authors would also like to thank Zhi Ye for the suggestions on benchmarking the computational cost of the model.
[118] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed deeponets, 2021.
[119] Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks, 2022.
[120] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
[121] Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. U-fno—an enhanced fourier neural operator-based deep-learning model for multiphase flow. Advances in Water Resources, 163:104180, 2022.
[122] Matthew A Wright and Joseph E Gonzalez. Transformers are deep infinite-dimensional non-mercer binary kernel machines. arXiv preprint arXiv:2106.01506, 2021.
[123] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystromformer: A nystrom-based algorithm for approximating self-attention, 2021.
[124] Yinchong Yang, Denis Krompass, and Volker Tresp. Tensor-train recurrent neural networks for video classification, 2017.
[125] Fisher Yu and Vladlen Koltun. Multi-scale context aggregation by dilated convolutions, 2016.
[126] Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontonon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences, 2021.
[127] Min Zhu, Handi Zhang, Anran Jiao, George Em Karniadakis, and Lu Lu. Reliable extrapolation of deep neural operators informed by physics 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.
Hyperparameter	2D Kolmogorov	3D turbulence	3D smoke	2D Darcy flow
Hidden dimension	128	128	128	128
Depth	4	4	3	3
Heads	8	6	6	12
Kernel dimension	128	192	192	128
Input encoder	2D Conv	2D Conv	2D Conv	MLP
Output decoder	MLP	MLP	MLP	MLP
Table 5: Major hyperparameters for FactFormer
Hidden dimension indicates the number of channels in the latent space. Depth denotes the number of attention layers. Before entering each attention layer, we modulate the latent encoding with positional encoding:  $z_{i} \gets z_{i} + \psi(x_{i})$ , where  $x_{i}$  is the Cartesian coordinate of latent encoding  $z_{i}$ ,  $\psi : \mathbb{R}^{n} \mapsto \mathbb{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:  $\kappa(x, \xi) = \sum_{l}^{d_{k}} q_{l}(x) k_{l}(\xi)$ . We train the model with AdamW optimizer and cyclic learning rate scheduler with a maximum learning rate  $3e - 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.