SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zijie Li, Dule Shu, Amir Barati Farimani
Mechanical Engineering Department
{zijieli, dules}@andrew.cmu.edu & barati@cmu.edu
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 \times 64 \times 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.
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
Figure 1: Model's prediction (pred.) and reference ground truth (ref.). Left: 2D Kolmogorov flow on  $256 \times 256$  grid; Right: 3D smoke buoyancy on  $64 \times 64 \times 64$  grid ( $zOy$  cross-section is shown).
PDE modeling [16, 35, 43, 54, 67, 74, 86]. In these works, attention is typically applied to every grid point in the domain to exploit both the local and non-local structure of the system, and therefore a linear-complexity variant of attention is usually necessary. As the number of grid points grows exponentially with respect 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.<sup>1</sup>
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 input-target 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
$^{1}$ Code for this project is available at: https://github.com/BaratiLab/FactFormer.
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.1 Attention mechanism
Standard attention Given three sets of vectors, namely the queries  $\{\mathbf{q}_i\}_{i=1}^{N_q}$ , keys  $\{\mathbf{k}_i\}_{i=1}^{N_k}$ , and values  $\{\mathbf{v}_i\}_{i=1}^{N_v}$  (assuming  $N_k = N_v$ ), attention mechanism [2, 33, 79, 115] dynamically computes a weighted average of the values:  $\mathbf{z}_i = \sum_{j=1}^{N_v} h(\mathbf{q}_i, \mathbf{k}_j) \mathbf{v}_j$ , where  $\mathbf{q}_i, \mathbf{k}_i, \mathbf{v}_i \in \mathbb{R}^{1 \times d}$ ,  $h(\cdot)$  is the weight function that determines the contribution of a specific value to the final output. An example of  $h(\cdot)$  is the scaled-dot product with softmax [115]:  $h(\mathbf{q}_i, \mathbf{k}_j) = \exp(\mathbf{q}_i \mathbf{k}_j^T / \tau) / \sum_s \exp(\mathbf{q}_i \mathbf{k}_s^T / \tau)$ , and  $\tau$  is usually chosen as  $\tau = \sqrt{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:
$$
\mathbf {q} _ {i} = \mathbf {u} _ {i} W _ {q}, \mathbf {k} _ {i} = \mathbf {u} _ {i} W _ {k}, \mathbf {v} _ {i} = \mathbf {u} _ {i} W _ {v}, \tag {1}
$$
where  $\mathbf{u}_i\in \mathbb{R}^{1\times d_{\mathrm{in}}}$  is the input vector and  $\{W_q,W_k,W_v\} \in \mathbb{R}^{d_{\mathrm{in}}\times 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\} \in \mathbb{R}^{N\times d}$ , their row vectors:  $\mathbf{q}_i / \mathbf{k}_i / \mathbf{v}_i$ , correspond to the sampling of a set of functions  $\{q_{l}(\cdot),k_{l}(\cdot),v_{l}(\cdot)\}_{l = 1}^{d}$  on grid point  $x_{i}$ , where  $\{x_{i}\}_{i = 1}^{N}$  discretizes the underlying domain. As a more concrete example, the  $l$ -th column (channel) of  $\mathbf{q}_i$ , represents the sampling of function  $q_{l}(\cdot)$  on a grid point, i.e.  $(\mathbf{q}_i)^l = q_l(x_i)$ . Furthermore, softmax-free attention is equivalent to the numerical quadrature of a kernel integral:
$$
\left(\mathbf {z} _ {i}\right) ^ {l} = \sum_ {s = 1} ^ {N} w _ {s} \left(\mathbf {q} _ {i} \cdot \mathbf {k} _ {s}\right) \left(\mathbf {v} _ {s}\right) ^ {l} \approx \int_ {\Omega} \kappa (x _ {i}, \xi) v _ {l} (\xi) d \xi , \tag {2}
$$
where  $\mathbf{z}_i$  is the output vector,  $\kappa(x,\xi) = \sum_{l=1}^{d} q_l(x) k_l(\xi)$  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},\xi)$ . 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  $\mathbf{q}_i,\mathbf{k}_j$  and their corresponding spatial coordinates  $x_{i},x_{j}$ , RoPE [105]  $(g(\cdot ,\cdot):\mathbb{R}^{1\times d}\times \mathbb{R}\mapsto \mathbb{R}^{1\times d})$  is defined as:
$$
g \left(\mathbf {q} _ {i}, x _ {i}\right) = \mathbf {q} _ {i} \boldsymbol {\Theta} \left(x _ {i}\right), \quad g \left(\mathbf {k} _ {j}, x _ {j}\right) = \mathbf {k} _ {j} \boldsymbol {\Theta} \left(x _ {j}\right) \tag {3}
$$
where:  $\Theta (x_{i}) = \mathrm{Diag}\left(R_{1}(x_{i}),\ldots ,R_{d / 2}(x_{i})\right),\quad R_{l} = \left[ \begin{array}{cc}\cos \left(\lambda x_{i}\theta_{l}\right) & -\sin \left(\lambda x_{i}\theta_{l}\right)\\ \sin \left(\lambda x_{i}\theta_{l}\right) & \cos \left(\lambda x_{i}\theta_{l}\right) \end{array} \right],$
and  $\lambda, \theta_{l}$  are hyperparameters.  $\theta_{l}$  is usually chosen as  $10000^{-2(l - 1) / d}, l \in \{1, 2, \dots, d / 2\}$  following Vaswani et al. [115] and Su et al. [105].  $\lambda$  is a mesh-based weight that we heuristically set to 64 throughout most problems. The projection function  $\Theta(\cdot): \mathbb{R} \mapsto \mathbb{R}^d \times \mathbb{R}^d$  can explicitly modulate the dot product with relative position:  $g(\mathbf{q}_i, x_i) g(\mathbf{k}_j, x_j)^T = \mathbf{q}_i \Theta(x_i - x_j) \mathbf{k}_j^T$ , 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):
$$
Z = w \tilde {Q} \tilde {K} ^ {T} V, \tag {4}
$$
where  $\tilde{\square}$  denotes a matrix whose row vectors are RoPE encoded as in (3), e.g.,  $\tilde{Q}_i = g(\mathbf{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  $\tilde{Q} \tilde{K}^T$  evaluates the kernel function  $\kappa(\cdot, \cdot)$  on the discretization grid  $\{x_i\}_{i=1}^N$ .
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  $\tilde{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} \times S_{2} \times \ldots \times 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
Figure 2: Schematic of the factorized kernel attention. Upper path: the input is transformed into the Value via a linear transformation. Lower path: the input is first projected into multiple sub-functions with a one-dimensional domain. These sub-functions are then used to derive the Query and Key on each axis, and their dot products form the kernel function of the corresponding axis. The Value is iteratively updated by the kernel integral transform along each axis and finally sent to an MLP.
$U \in \mathbb{R}^{S_1 \times S_2 \times \ldots \times S_n^2}$ . The product of it with a matrix  $W \in \mathbb{R}^{J \times S_m}$  across the  $m$ -th mode will result in a tensor of shape  $S_1 \times \ldots \times S_{m-1} \times J \times S_{m+1} \times \ldots \times S_n$ , whose elements are defined as:
$$
\left(U \times_ {m} W\right) _ {i _ {1} i _ {2} \dots i _ {m - 1} j i _ {m + 1} \dots i _ {n}} = \sum_ {i _ {m} = 1} ^ {S _ {m}} U _ {i _ {1} i _ {2} \dots i _ {m} \dots i _ {n}} W _ {j i _ {m}}. \tag {5}
$$
Learnable projection The first major component of the proposed framework is a set of learnable integral operators  $\{\mathcal{G}^{(1)},\mathcal{G}^{(2)},\ldots ,\mathcal{G}^{(n)}\}$  that projects the input function  $u:\mathbb{R}^n\mapsto \mathbb{R}^d$  into a set of functions with one-dimensional domain  $\{\phi^{(1)},\phi^{(2)},\dots,\phi^{(n)}\} \in \mathbb{R}\mapsto \mathbb{R}^d$ , which is defined as:
$$
\begin{array}{l} \phi^ {(m)} \left(x _ {i} ^ {(m)}\right) = \mathcal {G} ^ {(m)} (u) \left(x _ {i} ^ {(m)}\right) \tag {6} \\ = h ^ {(m)} \left(w \int_ {\Omega_ {1}} \dots \int_ {\Omega_ {n}} \gamma^ {(m)} \left(u \left(\xi_ {1}, \ldots , \xi_ {m - 1}, x _ {i} ^ {(m)}, \xi_ {m + 1}, \ldots , \xi_ {n}\right)\right) d \xi_ {1} \ldots d \xi_ {m - 1} d \xi_ {m + 1} \ldots d \xi_ {n}\right), \\ \end{array}
$$
where  $h^{(m)}(\cdot):\mathbb{R}^d\mapsto \mathbb{R}^d$  and  $\gamma^{(m)}(\cdot):\mathbb{R}^d\mapsto \mathbb{R}^d$  are pointwise learnable functions and  $w = 1 / (L_1L_2\dots L_{m - 1}L_{m + 1}\dots L_n)$ , with  $L_{m}$  being the size of domain  $\Omega_{m}$  discretized by  $\{x_i^{(m)}\}_{i = 1}^{S_m}$ . In practice, we implement  $h^{(m)}$  as a three-layer multi-layer perception (MLP) similar to the feedforward network in Transformer [115] and  $\gamma^{(m)}$  as a simple linear transformation. When the underlying grid is uniform, (6) simply amounts to first transforming the input with pointwise learnable functions  $\gamma^{(m)}(\cdot)$ , 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):
$$
\begin{array}{l} z \left(x _ {i _ {1}} ^ {(1)}, x _ {i _ {2}} ^ {(2)}, \ldots , x _ {i _ {n}} ^ {(n)}\right) \\ = \int_ {\Omega_ {1}} \kappa^ {(1)} \left(x _ {i _ {1}} ^ {(1)}, \xi_ {1}\right) \int_ {\Omega_ {2}} \kappa^ {(2)} \left(x _ {i _ {2}} ^ {(2)}, \xi_ {2}\right) \dots \int_ {\Omega_ {n}} \kappa^ {(n)} \left(x _ {i _ {n}} ^ {(n)}, \xi_ {n}\right) v \left(\xi_ {1}, \xi_ {2}, \dots , \xi_ {n}\right) d \xi_ {1} d \xi_ {2} \dots d \xi_ {n}, \tag {7} \\ \end{array}
$$
where kernels  $\{\kappa^{(1)},\ldots ,\kappa^{(n)}\} :\mathbb{R}\times \mathbb{R}\mapsto \mathbb{R}$  are computed based on projected single-dimensional function along each axis,  $v(\cdot):\mathbb{R}^n\mapsto \mathbb{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  $\{\hat{U}^{(1)},\dots,\hat{U}^{(n)}\}$ $(\hat{U}^{(m)}\in \mathbb{R}^{S_m\times d})$  from the input  $U\in \mathbb{R}^{S_1\times \dots \times S_n\times d}$ , where the  $i_{m}$ -th row of  $\hat{U}$  is the evaluation of projected function  $\phi^{(m)}(\cdot)$  at  $x_{i_m}\colon \hat{U}_{i_m}^{(m)} = \phi^{(m)}(x_{i_m})$ . Then we apply linear transformation on them to obtain the query/key matrix just as standard attention:  $Q^{(m)} = \hat{U}^{(m)}W_{q}^{(m)},K^{(m)} = \hat{U}^{(m)}W_{k}^{(m)}$ , where  $\{W_q^{(m)},W_k^{(m)}\} \in \mathbb{R}^{d\times d}$  are learnable matrices. The query and key are used to compute the kernel matrix  $A^{(m)}\in \mathbb{R}^{S_m\times S_m}$ :
$$
A ^ {(m)} = w _ {m} \tilde {Q} ^ {(m)} (\tilde {K} ^ {(m)}) ^ {T}, \tag {8}
$$
In practice we often have an additional mode for the channel, resulting in a  $(n + 1)$ -way tensor.
where  $w_{m}$  is the mesh weight,  $\tilde{\square}$  denotes the RoPE encoded matrix as discussed in (4), the  $i$ -th row and  $j$ -th column of  $A^{(m)}$  represents the kernel value  $\kappa^{(m)}(x_i^{(m)}, x_j^{(m)})$ . 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 \in \mathbb{R}^{S_1 \times \dots \times S_n \times d}$  is derived from the input via a linear transformation (i.e.  $(m + 1)$ -th mode product):  $V = U \times_{n+1} W_v$ , where  $W_v \in \mathbb{R}^{d \times d}$  is again a learnable matrix. The overall factorized kernel integral is numerically approximated with the following tensor-matrix product (Figure 2):