SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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_{\mathrm{in}})$  where  $N$  is the number of spatial grid points and  $T_{\mathrm{in}}$  is the number of input frames. Then we apply 2D convolution filters of size  $(1,T_{\mathrm{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
rollback 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  $5e - 4$  to  $5e - 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.
Figure 6: Error trend of vorticity  $\omega$  on 2D Kolmogorov flow. Dashed line: AR; Solid line: LM
Figure 7: Error trend of pressure  $p$  on 3D isotropic turbulence.
Figure 8: Error trend of velocity  $\mathbf{u}$  on 3D isotropic turbulence.
Figure 9: Error trend of marker field  $d$  on 3D smoke buoyancy.
Figure 10: Error trend of velocity field  $\mathbf{u}$  on 3D smoke buoyancy.
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):
$$
U _ {1} = \operatorname {A t t} \left(U _ {0}\right), \quad U _ {2} = \sigma \left(U _ {1} W _ {1}\right) W _ {2}, \tag {11}
$$
where  $U_0 \in \mathbb{R}^{N \times d}$  is the input,  $W_1, W_2 \in \mathbb{R}^{d \times d}$  are learnable weights for the feedforward network. Given a loss function  $l(\cdot): \mathbb{R}^{N \times d} \mapsto \mathbb{R}$  (for example, mean squared error), define  $\tilde{U}_1 \coloneqq U_1W_1$ ,  $\frac{\partial l}{\partial U_2} \coloneqq h$ ,  $\frac{\partial \sigma(\tilde{U}_1)}{\partial U_1} \coloneqq g$ , the gradient for the weights and input in the feedforward network are ( $\odot$  denotes element-wise multiplication):
$$
\frac {\partial l}{\partial W _ {2}} = \sigma \left(\tilde {U} _ {1}\right) ^ {T} h \tag {12}
$$
$$
\frac {\partial l}{\partial W _ {1}} = U _ {1} ^ {T} \left(h W _ {2} ^ {T} \odot g\right) \tag {13}
$$
$$
\frac {\partial l}{\partial U _ {1}} = \left(h W _ {2} ^ {T} \odot g\right) W _ {1} ^ {T} \tag {14}
$$
For linear (softmax-free) dot-product attention:  $U_{1} = \mathrm{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_{k}^{T}U_{0}^{T}U_{0}W_{v}$ , the gradient of weights are:
$$
\frac {\partial l}{\partial W _ {q}} = U _ {0} ^ {T} \frac {\partial l}{\partial U _ {1}} W _ {v} ^ {T} U _ {0} ^ {T} U _ {0} W _ {k}, \tag {15}
$$
$$
\frac {\partial l}{\partial W _ {k}} = U _ {0} ^ {T} U _ {0} W _ {v} \left(\frac {\partial l}{\partial U _ {1}}\right) ^ {T} U _ {0} W _ {q}, \tag {16}
$$
$$
\frac {\partial l}{\partial W _ {v}} = U _ {0} ^ {T} U _ {0} W _ {k} W _ {q} ^ {T} U _ {0} ^ {T} \frac {\partial l}{\partial U _ {1}}, \tag {17}
$$
where  $U_0 \in \mathbb{R}^{N \times 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:
$$
Z _ {i, c} = \sum_ {m = 1} ^ {d} Q _ {i, m} \left(\sum_ {j = 1} ^ {N} K _ {j, m} V _ {j, c}\right), \tag {18}
$$
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 \in \mathbb{R}^{N \times d}$  is input. The gradient of parameters are computed as:
$$
\frac {\partial Z _ {i , c}}{\partial \left(W _ {q}\right) _ {r , s}} = X _ {i, r} \left(\sum_ {j = 1} ^ {N} K _ {j, s} V _ {j, c}\right), \tag {19}
$$
$$
\frac {\partial Z _ {i , c}}{\partial \left(W _ {k}\right) _ {r , s}} = Q _ {i, s} \left(\sum_ {j = 1} ^ {N} X _ {j, r} V _ {j, c}\right), \tag {20}
$$
$$
\frac {\partial Z _ {i , c}}{\partial \left(W _ {v}\right) _ {r , s}} = \left\{ \begin{array}{l l} \sum_ {m = 1} ^ {d} Q _ {i, m} \left(\sum_ {j = 1} ^ {N} K _ {j, m} X _ {i, r}\right) & : \text {i f} s = c \\ 0 & : \text {o t h e r w i s e .} \end{array} \right. \tag {21}
$$
For the proposed factorized attention, it can written as:
$$
Z _ {i, c} = \sum_ {j _ {1} = 1} ^ {S _ {1}} \sum_ {j _ {2} = 1} ^ {S _ {2}} \dots \sum_ {j _ {n} = 1} ^ {S _ {n}} A _ {i _ {1}, j _ {1}} ^ {(1)} A _ {i _ {2}, j _ {2}} ^ {(2)} \dots A _ {i _ {n}, j _ {n}} ^ {(n)} V _ {j, c}, \quad j := (j _ {1}, j _ {2}, \dots , j _ {n}), i := i _ {1}, i _ {2}, \dots , i _ {n} \tag {22}
$$
where  $A^{(m)} = Q^{(n)}(K^{(m)})^T$  ( $A^{(m)} \in \mathbb{R}^{S_m \times S_m}$ ,  $N = S_1 \times S_2 \times \ldots \times S_n$ ) is the axial kernel matrix as defined in (8),  $X^{(m)} \in \mathbb{R}^{S_m \times d}$  is the axial projection along  $m$ -th axis as defined in (6).
For  $W_{v}$ , its gradient is computed as:
$$
\frac {\partial Z _ {i , c}}{\partial \left(W _ {v}\right) _ {r , s}} = \left\{ \begin{array}{l l} \sum_ {j _ {1} = 1} ^ {S _ {1}} \sum_ {j _ {2} = 1} ^ {S _ {2}} \dots \sum_ {j _ {n} = 1} ^ {S _ {n}} A _ {i _ {1}, j _ {1}} ^ {(1)} A _ {i _ {2}, j _ {2}} ^ {(2)} \dots A _ {i _ {n}, j _ {n}} ^ {(n)} X _ {j, r} & : \text {i f} s = c \\ 0 & : \text {o t h e r w i s e .} \end{array} \right. \tag {23}
$$
For  $W_{q}^{(n)}, W_{k}^{(n)}$ , their gradients are computed as:
$$
\frac {\partial Z _ {i , c}}{\partial \left(W _ {q} ^ {(n)}\right) _ {r , s}} = X _ {i _ {n}, r} ^ {(n)} \sum_ {j _ {1} = 1} ^ {S _ {1}} \sum_ {j _ {2} = 1} ^ {S _ {2}} \dots \sum_ {j _ {n} = 1} ^ {S _ {n}} A _ {i _ {1}, j _ {1}} ^ {(1)} A _ {i _ {2}, j _ {2}} ^ {(2)} \dots A _ {i _ {n - 1}, j _ {n - 1}} ^ {(n - 1)} V _ {j, c} K _ {j _ {n}, s} ^ {(n)}, \tag {24}
$$
$$
\frac {\partial Z _ {i , c}}{\partial \left(W _ {k} ^ {(n)}\right) _ {r , s}} = Q _ {i _ {n, s}} ^ {(n)} \sum_ {j _ {1} = 1} ^ {S _ {1}} \sum_ {j _ {2} = 1} ^ {S _ {2}} \dots \sum_ {j _ {n} = 1} ^ {S _ {n}} A _ {i _ {1}, j _ {1}} ^ {(1)} A _ {i _ {2}, j _ {2}} ^ {(2)} \dots A _ {i _ {n - 1}, j _ {n - 1}} ^ {(n - 1)} V _ {j, c} X _ {j _ {n}, r} ^ {(n)}, \tag {25}
$$
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.
(a) Forward  $^+$  backward peak memory usage.
(b) Forward  $^+$  backward time per iteration.
Figure 11: Benchmark of factorized attention and linear attention on 2D domain (with  $128 \times 128$  grid) with varying kernel dimension. Benchmark is done on an RTX 3090 with PyTorch 1.8.2 and a batch size of 4. Hyperparameter setting is the same as in Table 5-2D Kolmogorov flow. "OOM" denotes out of memory.
(a) Forward  $^+$  backward peak memory usage.
Figure 12: Benchmark of factorized attention and linear attention on 3D domain with varying grid size. Benchmark is done on an A6000 with PyTorch 1.8.2 and batch size of 1. Hyperparameter setting is the same as in Table 5-3D isotropic turbulence.
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 \times 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.
Figure 13: Ablation study on the key hyperparameters. Red color denotes the final choice of hyperparameter. Experiments are carried out on the validation fold.
(b) Ablation on the number of heads.
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  $\kappa (\xi_1,\xi_2)$  depends only on the relative distance between two points, e.g.  $L^2$  distance:  $\lvert\lvert\xi_1-\xi_2\rvert\rvert_2\rangle$ ). 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
Figure 14: Visualization of normalized attention kernel.
(d) Layer 4, Head 1: kernel of  $x$  axis
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.
Figure 15: Averaged frame-wise loss trends. Each loss curve corresponds to a model initialized under a specific seed.
In this section we provide the details for each dataset.
2D Kolmogorov flow The incompressible Navier-Stokes equation under vorticity form reads as,
$$
\begin{array}{l} \frac {\partial \omega (\mathbf {x} , t)}{\partial t} + \mathbf {u} (\mathbf {x}, t) \cdot \nabla \omega (\mathbf {x}, t) = \frac {1}{R e} \nabla^ {2} \omega (\mathbf {x}, t) + f (\mathbf {x}), \quad \mathbf {x} \in (0, 2 \pi) ^ {2}, t \in (0, T ], \\ \nabla \cdot \mathbf {u} (\mathbf {x}, t) = 0, \quad \mathbf {x} \in (0, 2 \pi) ^ {2}, t \in [ 0, T ], \tag {26} \\ \omega (\mathbf {x}, 0) = \omega_ {0} (\mathbf {x}), \quad \mathbf {x} \in (0, 2 \pi) ^ {2}, \\ \end{array}
$$
where  $\omega$  denotes vorticity,  $\mathbf{u}$  denotes velocity,  $Re$  denotes the Reynolds number,  $\mathbf{x} = (x_{1}, x_{2})$  denotes the spatial coordinates and  $f(\cdot)$  is the forcing term that is set to  $f(\mathbf{x}) = -n\cos(nx_{2}) - 0.1\omega(\mathbf{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\omega(\mathbf{x})$  as described in Kochkov et al. [56]. The initial condition  $\omega_{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 \times 2048$  and a temporal resolution of  $1e - 4$ .
3D isotropic turbulence The incompressible Navier-Stokes equation for this problem is given as:
$$
\begin{array}{l} \frac {\partial \mathbf {u} (\mathbf {x} , t)}{\partial t} + \mathbf {u} (\mathbf {x}, t) \cdot \nabla \mathbf {u} (\mathbf {x}, t) = \nu \nabla^ {2} \mathbf {u} (\mathbf {x}, t) - \frac {1}{\rho} \nabla p (\mathbf {x}, t) + \mathbf {f} (\mathbf {x}), \quad \mathbf {x} \in (0, 2 \pi) ^ {3}, t \in (0, T ], \\ \nabla \cdot \mathbf {u} (\mathbf {x}, t) = 0, \quad \mathbf {x} \in (0, 2 \pi) ^ {3}, t \in [ 0, T ], \tag {27} \\ \mathbf {u} (\mathbf {x}, 0) = \mathbf {u} _ {0} (\mathbf {x}), \quad \mathbf {x} \in (0, 2 \pi) ^ {3}, \\ \end{array}
$$
where  $\mathbf{u}$  denotes velocity,  $p$  denotes the pressure,  $\nu$  is the viscosity parameter,  $\mathbf{x} = [x_1,x_2,x_3]$  denotes the spatial coordinates and  $f(\cdot)$  is the forcing term. The equation is periodic in all three spatial dimensions. The initialization of  $\mathbf{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 \times 60 \times 60$  and a temporal resolution of  $0.005s$ .
3D smoke buoyancy The governing equations for the 3D smoke buoyancy problem are incompressible Navier-Stokes equation (similar as above) coupled with advection equation:
$$
\begin{array}{l} \frac {\partial \mathbf {u} (\mathbf {x} , t)}{\partial t} + \mathbf {u} (\mathbf {x}, t) \cdot \nabla \mathbf {u} (\mathbf {x}, t) = \nu \nabla^ {2} \mathbf {u} (\mathbf {x}, t) - \frac {1}{\rho} \nabla p (\mathbf {x}, t) + \mathbf {f} (\mathbf {x}, t), \quad \mathbf {x} \in (0, L) ^ {3}, t \in (0, T ], \\ \frac {\partial d (\mathbf {x} , t)}{\partial t} + \mathbf {u} (\mathbf {x}, t) \cdot \nabla d (\mathbf {x}, t) = 0, \quad \mathbf {x} \in (0, L) ^ {3}, t \in (0, T ], \\ \nabla \cdot \mathbf {u} (\mathbf {x}, t) = 0, \quad \mathbf {x} \in (0, L) ^ {3}, t \in [ 0, T ], \\ \mathbf {u} (\mathbf {x}, 0) = 0, \quad \mathbf {d} (\mathbf {x}, 0) = d _ {0} (\mathbf {x}), \quad \mathbf {x} \in (0, L) ^ {3}, \\ \end{array}
$$
where  $\mathbf{f}(\mathbf{x},t) = [0,0,\eta d(\mathbf{x},t)]$ ,  $\eta$  is the buoyancy factor, the velocity field  $\mathbf{u}$  has a Dirichlet boundary condition:  $\mathbf{u}(\mathbf{x},\cdot) = 0$ ,  $\forall \mathbf{x}\in \partial \Omega$ , and the scalar density field for smoke has a Neumann boundary condition:  $\nabla d(\mathbf{x},\cdot) = 0$ ,  $\forall \mathbf{x}\in \partial \Omega$ . The initial condition  $d_0(\mathbf{x})$  is a random field 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 \times 64 \times 64$  and a temporal resolution of  $0.75s$ .
2D Darcy flow The equation for the 2D Darcy flow is defined as:
$$
- \nabla \cdot (a (x) \nabla u (x)) = f (x), \quad x \in (0, 1) ^ {2}, \tag {28}
$$
$$
u _ {0} (x) = 0, \quad x \in \partial (0, 1) ^ {2},
$$
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 \times 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.
3 Implemented with phiflow's Noise class, see:https://tum-pbs.github.io/PhiFlow/phi/field/
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.
Figure 16: Sample 1 of 2D Kolmogorov flow.
Figure 17: Sample 2 of 2D Kolmogorov flow.
Figure 18: Pressure in 3D isotropic turbulence.
Figure 19:  $x$ -component of velocity in 3D isotropic turbulence.
Figure 20:  $y$ -component of velocity in 3D isotropic turbulence.
Figure 21:  $z$ -component of velocity in 3D isotropic turbulence.
Figure 22: Smoke marker field in 3D smoke buoyancy.
Figure 23:  $x$ -component of velocity in 3D smoke buoyancy.
Figure 24:  $y$ -component of velocity in 3D smoke buoyancy.
Figure 25:  $z$ -component of velocity in 3D smoke buoyancy.
Figure 26: Flow field of 2D Darcy flow.
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
Figure 27: An illustrative example of FactFormer and Axial Transformer applying to 2D input data, with some details such as positional encoding, multi-head mechanism and softmax in Axial Transformer omitted for simplicity. For Axial Transformer, column-wise attention block is shown as an example.