Choose a Transformer: Fourier or Galerkin
Year: 2021

Abstract:
In this paper, we apply the self-attention from the state-of-the-art Transformer in Attention Is All You Need [88] for the first time to a data-driven operator learning problem related to partial differential equations. An effort is put together to explain the heuristics of, and to improve the efficacy of the attention mechanism. By employing the operator approximation theory in Hilbert spaces, it is demonstrated for the first time that the softmax normalization in the scaled dot-product attention is sufficient but not necessary. Without softmax, the approximation capacity of a linearized Transformer variant can be proved to be comparable to a Petrov-Galerkin projection layer-wise, and the estimate is independent with respect to the sequence length. A new layer normalization scheme mimicking the Petrov-Galerkin projection is proposed to allow a scaling to propagate through attention layers, which helps the model achieve remarkable accuracy in operator learning tasks with unnormalized data. Finally, we present three operator learning experiments, including the viscid Burgers’ equation, an interface Darcy flow, and an inverse interface coefficient identification problem. The newly proposed simple attention-based operator learner, Galerkin Transformer, shows significant improvements in both training cost and evaluation accuracy over its softmax-normalized counterparts.

heading: 1 Introduction
text: Partial differential equations (PDEs) arise from almost every multiphysics and biological systems, from the interaction of atoms to the merge of galaxies, from the formation of cells to the change of climate. Scientists and engineers have been working on approximating the governing PDEs of these physical systems for centuries. The emergence of the computer-aided simulation facilitates a cost-friendly way to study these challenging problems. Traditional methods, such as finite element/difference [20, 22], spectral methods [12], etc., leverage a discrete structure to reduce an infinite dimensional operator map to a finite dimensional approximation problem. Meanwhile, in the field practice of many scientific disciplines, substantial data for PDE-governed phenomena available on discrete grids enable modern black-box models like Physics-Informed Neural Network (PINN) [71, 62, 49] to exploit measurements on collocation points to approximate PDE solutions.
Nonetheless, for traditional methods or data-driven function learners such as PINN, given a PDE, the focus is to approximate a single instance, for example, solving for an approximated solution for one coefficient with a fixed boundary condition. A slight change to this coefficient invokes a potentially expensive re-training of any data-driven function learners. In contrast, an operator learner aims to learn a map between infinite-dimensional function spaces, which is much more difficult yet rewarding. A well-trained operator learner can evaluate many instances without re-training or collocation points, thus saving valuable resources, and poses itself as a more efficient approach in the long run. Data-driven resolution-invariant operator learning is a booming new research direction [60, 5, 56, 64, 90, 57, 61, 91, 37, 74]. The pioneering model, DeepONet [60], attributes architecturally to a universal approximation theorem for operators [18]. Fourier Neural Operator (FNO) [57] notably
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
shows an awing state-of-the-art performance outclassing classic models such as the one in [100] by orders of magnitudes in certain benchmarks.
Under a supervised setting, an operator learner is trained with the operator’s input functions and their responses to the inputs as targets. Since both functions are sampled at discrete grid points, this is a special case of a seq2seq problem [81]. The current state-of-the-art seq2seq model is the Transformer first introduced in [88]. As the heart and soul of the Transformer, the scaled dot-product attention mechanism is capable of unearthing the hidden structure of an operator by capturing long-range interactions. Inspired by many insightful pioneering work in Transformers [50, 19, 75, 84, 96, 97, 95, 59, 76, 66], we have modified the attention mechanism minimally yet in a mathematically profound manner to better serve the purpose of operator learning.
Among our new Hilbert space-inspired adaptations of the scaled dot-product attention, the first and foremost change is: no softmax, or the approximation thereof. In the vanilla attention [88], the softmax succeeding the matrix multiplication convexifies the weights for combining different positions’ latent representations, which is regarded as an indispensable ingredient in the positive kernel interpretation of the attention mechanism [84]. However, softmax acts globally in the sequence length dimension for each row of the attention matrix, and further adds to the quadratic complexity of the attention in the classic Transformer. Theory-wise, instead of viewing “row ≈ word” in the Natural Language Processing (NLP) tradition, the columns of the query/keys/values are seen as sampling of functions in Hilbert spaces on discretized grids. Thus, taking the softmax away allows us to verify a discrete Ladyzhenskaya–Babuška–Brezzi (LBB) condition, which further amounts to the proof that the newly proposed Galerkin-type attention can explicitly represent a Petrov-Galerkin projection, and this approximation capacity is independent of the sequence length (Theorem 4.3).
Numerically, the softmax-free models save valuable computational resources, outperforming the ones with the softmax in terms of training FLOP and memory consumption (Section 5). Yet in an ablation study, the training becomes unstable for softmax-free models (Table 8). To remedy this, a new Galerkin projection-type layer normalization scheme is proposed to act as a cheap diagonal alternative to the normalizations explicitly derived in the proof of the Petrov-Galerkin interpretation (equation (40)). Since a learnable scaling can now be propagated through the encoder layers, the attentionbased operator learner with this new layer normalization scheme exhibits better comprehension of certain physical properties associated with the PDEs such as the energy decay. Combining with other approximation theory-inspired tricks including a diagonally dominant rescaled initialization for the projection matrices and a layer-wise enrichment of the positional encodings, the evaluation accuracies in various operator learning tasks are boosted by a significant amount.
Main contributions. The main contributions of this work are summarized as follows.
• Attention without softmax. We propose a new simple self-attention operator and its linear variant without the softmax normalization. Two new interpretations are offered, together with the approximation capacity of the linear variant proved comparable to a Petrov-Galerkin projection.
• Operator learner for PDEs. We combine the newly proposed attention operators with the current best state-of-the-art operator learner Fourier Neural Operator (FNO) [57] to significantly improve its evaluation accuracy in PDE solution operator learning benchmark problems. Moreover, the new model is capable of recovering coefficients based on noisy measurements that traditional methods or FNO cannot accomplish.
• Experimental results. We present three benchmark problems to show that operator learners using the newly proposed attentions are superior in computational/memory efficiency, as well as in accuracy versus those with the conventional softmax normalization. The PyTorch codes to reproduce our results are available as an open-source software. 1

heading: 2 Related Works
text: Operator learners related to PDEs. In [4, 5], certain kernel forms of the solution operator of parametric PDEs are approximated using graph neural networks. The other concurrent notable approach is DeepONet [60, 61]. [56] further improves the kernel approach by exploiting the multilevel grid structure. [57] proposes a discretization-invariant operator learner to achieve a state-of-the-art
1https://github.com/scaomath/galerkin-transformer
performance in certain benchmark problems. [90, 91] proposed a DeepONet roughly equivalent to an additive attention, similar to the one in the Neural Turing Machine (NMT) in [7]. Model/dimension reduction combined with neural nets is another popular approach to learn the solution operator for parametric PDEs [10, 64, 55, 24]. Deep convolutional neural networks (DCNN) are widely applied to learn the solution maps with a fixed discretization size [1, 9, 40, 36, 35, 100, 86]. Recently, DCNN has been successfully applied in various inverse problems [35, 47] such as Electrical Impedance Tomography (EIT). To our best knowledge, there is no work on data-driven approaches to an inverse interface coefficient identification for a class of coefficients with random interface geometries.
Attention mechanism and variants. Aside from the ground-breaking scaled dot-product attention in [88], earlier [7] proposed an additive content-based attention, however, with a vanishing gradient problem due to multiple nonlinearity composition. [25] shows the first effort in removing the softmax normalization in [7] after the projection, however, it still uses a Sigmoid nonlinearity before the additive interpolation propagation stage, and performs worse than its softmax counterpart. The current prevailing approach to linearize the attention leverages the assumption of the existence of a feature map to approximate the softmax kernel [50, 19, 70]. Another type of linearization exploits the low-rank nature of the matrix product using various methods such as sampling or projection [73, 11, 79, 92], or fast multipole decomposition [65]. The conjecture in [75] inspires us to remove the softmax overall. [76] first proposed the inverse sequence length scaling normalization for a linear complexity attention without the softmax, however, the scaling normalization has not been extensively studied in examples and performs worse.
Various studies on Transformers. The kernel interpretation in [84] inspires us to reformulate the attention using the Galerkin projection. [95, Theorem 2] gives a theoretical foundation of removing the softmax normalization to formulate the Fourier-type attention. The Nyström approximation [97] essentially acknowledges the similarity between the attention matrix and an integral kernel. [96, 66, 59] inspires us to try different layer normalization and the rescaled diagonally dominant initialization schemes. The practices of enriching the latent representations with the positional encoding recurrently in our work trace back to [2, 26], and more recently, contribute to the success of AlphaFold 2 [48], as it is rewarding to exploit the universal approximation if the target has a dependence ansatz in the coordinate frame and/or transformation group but hard to be explicitly quantified. Other studies on adapting the attention mechanisms to conserve important physical properties are in [82, 31, 44].

heading: 3 Operator learning related to PDEs
text: Closely following the setup in [56, 57], we consider a data-driven model to approximate a denselydefined operator T : H1 → H2 between two Hilbert spaces with an underlying bounded spacial domain Ω ⊂ Rm. The operator T to be learned is usually related to certain physical problems, of which the formulation is to seek the solution to a PDE of the following two types.
Parametric PDE: given coefficient a ∈ A, and source f ∈ Y , find u ∈ X such that La(u) = f . (i) To approximate the nonlinear mapping from the varying parameter a to the solution with a
fixed right-hand side, T : A → X , a 7→ u. (ii) The inverse coefficient identification problem to recover the coefficient from a noisy measure-
ment ũ of the steady-state solution u, in this case, T : X → A, ũ 7→ a. Nonlinear initial value problem: given u0 ∈ H0, find u ∈ C([0, T ];H) such that ∂tu+N(u) = 0.
(iii) Direct inference from the initial condition to the solution. T : H0 → H, u0(·) 7→ u(t1, ·) with t1 ∆t with t1 much greater than the step-size in traditional explicit integrator schemes.
Using (i) as an example, based on the given N observations {a(j), u(j)}Nj=1 and their approximations {a(j)h , u (j) h } defined at a discrete grid of size h 1, the goal of our operator learning problem is to build an approximation Tθ to T , such that Tθ(ah) is a good approximation to u = L−1a f =: T (a) ≈ uh independent of the mesh size h, where ah and uh are in finite dimensional spaces Ah,Xh on this grid. We further assume that a(j) ∼ ν for a measure ν compactly supported on A, and the sampled data form a reasonably sized subset of A representative of field applications. The loss J (θ) is
J (θ) := Ea∼ν [ ‖ ( Tθ(a)− u‖2H + G(a, u; θ) ] (1)
and in practice is approximated using the sampled observations on a discrete grid
J (θ) ≈ 1 N N∑ j=1 {∥∥(Tθ(a(j)h )− u(j)h ∥∥2H + G(a(j)h , u(j)h ; θ)}. (2) In example (i), ‖·‖H is the standard L2-norm, and G(a, u; θ) serves as a regularizer with strength γ and is problem-dependent. In Darcy flow where La := −∇ · (a∇(·)), it is γ‖a∇(Tθ(a)− u)‖2L2(Ω), since u ∈ H1+α(Ω) (α > 0 depends on the regularity of a) and a∇u ∈H(div; Ω) a priori. For the evaluation metric, we drop the G(a, u; θ) term, and monitor the minimization of (2) using ‖ · ‖H.

heading: 4 Attention-based operator learner
text: Feature extractor. We assume the functions in both inputs and targets are sampled on a uniform grid. In an operator learning problem on Ω ⊂ R1, a simple feedforward neural network (FFN) is used as the feature extractor that is shared by every position (grid point).
Interpolation-based CNN. If Ω ⊂ R2, inspired by the multilevel graph kernel network in [56], we use two 3-level interpolation-based CNNs (CiNN) as the feature extractor, but also as the downsampling and upsampling layer, respectively, in which we refer to restrictions/prolongations between the coarse/fine grids both as interpolations. For the full details of the network structure please refer to Appendix B.
Recurrent enrichment of positional encoding. The Cartesian coordinates of the grid, on which the attention operator’s input latent representation reside, are concatenated as additional feature dimension(s) to the input, as well as to each latent representation in every attention head.
Problem-dependent decoder. The decoder is a problem-dependent admissible network that maps the learned representations from the encoder back to the target dimension. For smooth and regular solutions in H1+α(Ω), we opt for a 2-layer spectral convolution that is the core component in [57]. A simple pointwise feedforward neural network (FFN) is used for nonsmooth targets in L∞(Ω).

heading: 4.1 Simple self-attention encoder
text: The encoder contains a stack of identical simple attention-based encoder layers. For simplicity, we consider a single attention head that maps y ∈ Rn×d to another element in Rn×d, and define the
trainable projection matrices, and the latent representations Q/K/V as follows. WQ,WK ,WV ∈ Rd×d, and Q := yWQ, K := yWK , V := yWV . (3)
We propose the following simple attention that (i) uses a mesh (inverse sequence length)-weighted normalization without softmax, (ii) allows a scaling to propagate through the encoder layers. Attnsp : Rn×d → Rn×d, ỹ← y + Attn†(y), y 7→ ỹ + g(ỹ), (4) where the head-wise normalizations are applied pre-dot-product: for † ∈ {f, g},
(Fourier-type attention) z = Attnf(y) := (Q̃K̃>)V/n, (5)
(Galerkin-type attention) z = Attng(y) := Q(K̃>Ṽ )/n, (6) and ̃ denotes a trainable non-batch-based normalization. As in the classic Transformer [88], and inspired by the Galerkin projection interpretation, we choose ̃ as the layer normalization Ln( ), and g(·) as the standard 2-layer FFN identically applied on every position (grid point). In simple attentions, the weight for each row of V , or column of Q in the linear variant, is not all positive anymore. This can be viewed as a cheap alternative to the cosine similarity-based attention. Remark 4.1. If we apply the regular layer normalization rule that eliminates any scaling:
y 7→ Ln ( y + Attn†(y) + g ( Ln(y + Attn†(y)) )) , where Attn†(y) := Q(K>V )/n, (7)
then this reduces to the efficient attention first proposed in [76].

heading: 4.1.1 Structure-preserving feature map as a function of positional encodings
text: Consider an operator learning problem with an underlying domain Ω ⊂ R1. {xi}ni=1 denotes the set of grid points in the discretized Ω such that the weight 1/n = h is the mesh size. Let ζq(·), φk(·), ψv(·) : Ω → R1×d denote the feature maps of Q,K, V , i.e., the i-th row of Q,K, V written as qi = ζq(xi), ki = φk(xi), vi = ψv(xi). They are, in the NLP convention, viewed as the feature (embedding) vector at the i-th position, respectively. The inter-position topological structure such as continuity/differentiability in the same feature dimension is learned thus not explicit. The following ansatz for Q/K/V in the same attention head is fundamental to our new interpretations. Assumption 4.2. The columns of Q/K/V , respectively, contain the vector representations of the learned basis functions spanning certain subspaces of the latent representation Hilbert spaces.
Using V ∈ Rn×d with a full column rank as an example, its columns contain potentially a set of bases {vj(·)}dj=1 evaluated at the grid points (degrees of freedom, or DoFs). Similarly, the learned bases whose DoFs form the columns of Q,K are denoted as {qj(·)}dj=1, {kj(·)}dj=1, as well as {zj(·)}dj=1 for the outputs in (5) and (6). To be specific, the j-th column of V , denoted by vj , then stands for a vector representation of the j-th basis function evaluated at each grid point, i.e., its l-th position stands for (vj)l = vj(xl). Consequently, the row vi = (v1(xi), . . . , vd(xi)) can be alternatively viewed as the evaluation of a vector latent basis function at xi.
4.1.2 Fourier-type attention of a quadratic complexity
In the Fourier-type attention (5), Q,K are assumed to be normalized for simplicity, the j-th column (1 ≤ j ≤ d) in the i-th row (1 ≤ i ≤ n) of z is computed by (see Figure 2):
(zi)j = h(QK >)i• v j = h ( qi · k1, . . . , qi · kl, . . . , qi · kn )> · vj = h
n∑ l=1 (qi · kl)(vj)l ≈ ∫ Ω ( ζq(xi) · φk(ξ) ) vj(ξ) dξ,
(8)
where the h-weight facilitates the numerical quadrature interpretation of the inner product. Concatenating columns 1 ≤ j ≤ d yields the i-row zi of the output z: zi ≈ ∫ Ω ( ζq(xi) · φk(ξ) ) ψv(ξ) dξ. Therefore, without the softmax nonlinearity, the local dot-product attention output at i-th row computes approximately an integral transform with a non-symmetric learnable kernel function κ(x, ξ) := ζq(x)φk(ξ) evaluated at xi, whose approximation property has been studied in [95, Theorem 2], yet without the logits technicality due to the removal of the softmax normalization.
After the skip-connection, if we further exploit the learnable nature of the method and assume WV = diag{δ1, · · · , δd} such that δj 6= 0 for 1 ≤ j ≤ d, under Assumption 4.2: δ−1j vj(x) ≈ zj(x)− ∫
Ω
κ(x, ξ)vj(ξ) dξ, for j = 1, · · · , d, and x ∈ {xi}ni=1. (9)
This is the forward propagation of the Fredholm equation of the second-kind for each vj(·). When using an explicit orthogonal expansion such as Fourier to solve for {vj(·)}dj=1, or to seek for a better set of {vj(·)} in our case, it is long known being equivalent to the Nyström’s method with numerical integrations [8] (similar to the h = 1/n weighted sum). Therefore, the successes of the random Fourier features in [19, 70] and the Nyströmformer’s approximation [97] are not surprising.
Finally, we name this type of simple attention “Fourier” is due to the striking resemblance between the scaled dot-product attention and a Fourier-type kernel [30] integral transform, since eventually the target resides in a Hilbert space with an underlying spacial domain Ω, while the latent representation space parallels a “frequency” domain on Ω∗. This also bridges the structural similarity of the scaled dot-product attention with the Fourier Neural Operator [57] where the Fast Fourier Transform (FFT) can be viewed as a non-learnable change of basis.
4.1.3 Galerkin-type attention of a linear complexity
For the Galerkin-type simple attention in (6), K,V are assumed to be normalized for simplicity, we first consider the i-th entry in the j-th column zj of z (see Figure 3):
(zj)i = h q > i · (K>V )•j , (10)
which is the inner product of the i-th row of Q and the j-th column of K>V . Thus,
zj = h ( | | | | q1 q2 · · · qn | | | | )> (K>V )•j = h (K>V )>•j  q1...
qd > (11) This reads as: (K>V )•j contains the coefficients for the linear combination of the vector representations {ql}dl=1 of the bases stored in Q’s column space to form the output z. Meanwhile, the j-th column (K>V )•j of K>V consists the inner product of j-th column of V with every column of K.
zj = h d∑ l=1 ql(K>V )lj , where (K>V )•j = ( k1 · vj ,k2 · vj , · · · ,kd · vj )> . (12)
As a result, using Assumption 4.2, and for simplicity the latent Hilbert spaces Q,K,V are assumed to be defined on the same spacial domain Ω, i.e., kl(·), vj(·) evaluated at every xi are simply their vector representations kl (1 ≤ l ≤ d) and vj , we have the functions represented by the columns of the output z can be then compactly written as: rewriting 〈vj , kl〉 := (K>V )lj
zj(x) := d∑ l=1 〈vj , kl〉 ql(x), for j = 1, · · · , d, and x ∈ {xi}ni=1, (13)
where the bilinear form 〈·, ·〉 : V × K → R. (13) can be also written in a componentwise form:
zj(xi) := (z j)i = h d∑ l=1 (kl · vj)(ql)i ≈ d∑ l=1 (∫ Ω vj(ξ)kl(ξ) dξ ) ql(xi). (14)
Therefore, when { j(·)}dj=1, ∈ {q, k, v} consist approximations to three sets of bases for potentially different subspaces, and if we set the trial spaces as the column spaces of Q and the test space as that of K, respectively, the forward propagation of the Galerkin-type attention is a recast of a learnable Petrov–Galerkin-type projection (cf. Appendix D.1) for every basis represented by the columns of V . While the form of (14) suggests the orthonormality of the basis represented by Q,K, V , as well as being of full column ranks, the learnable nature of the method suggests otherwise (see Appendix D). At last, we have the following strikingly simple yet powerful approximation result. Theorem 4.3 (Céa-type lemma, simplified version). Consider a Hilbert space H defined on a bounded domain Ω ⊂ Rm discretized by n grid points, and f ∈ H. y ∈ Rn×d is the current latent representation for n > d > m and full column rank. Qh ⊂ Q ⊂ H and Vh ⊂ V ⊂ H are the latent approximation subspaces spanned by basis functions with the columns of Q and V in (3) as degrees of freedom, respectively, and 0 < dimQh = r ≤ dimVh = d. Let b(·, ·) : V ×Q → R be a continuous bilinear form, and if for any fixed q ∈ Qh the functional norm of b(·, q) is bounded below by c > 0, then there exists a learnable map gθ(·) that is the composition of the Galerkin-type attention operator with an updated set of projection matrices {WQ,WK ,WV }, and a pointwise universal approximator, such that for fh ∈ Qh being the best approximation of f in ‖ · ‖H it holds:
‖f − gθ(y)‖H ≤ c−1 min q∈Qh max v∈Vh |b(v, fh − q)| ‖v‖H + ‖f − fh‖H. (15)
Remarks on and interpretations of the best approximation result. Theorem 4.3 states that the Galerkin-type attention has the architectural capacity to represent a quasi-optimal approximation in ‖ · ‖H in the current subspace Qh. For the mathematically rigorous complete set of notations and the full details of the proof we refer the readers to Appendix D.3. Even though Theorem 4.3 is presented for a single instance of f ∈ H for simplicity, the proof shows that the attention operator is fully capable of simultaneously approximating a collection of functions (Appendix D.3.4).
Estimate (15) comes with great scalability with respect to the sequence length in that it all boils down to whether c is independent of n in the lower bound of ‖b(·, q)‖V′h . The existence of an n-independent lower bound is commonly known as the discrete version of the Ladyzhenskaya–Babuška–Brezzi (LBB) condition [21, Chapter 6.12], also referred as the Banach-Nečas-Babuška (BNB) condition in Galerkin methods on Banach spaces [29, Theorem 2.6].
As the cornerstone of the approximation to many PDEs, the discrete LBB condition establishes the surjectivity of a map from Vh to Qh. In a simplified context (15) above of approximating functions using this linear attention variant (Q: values, query, V : keys), it roughly translates to: for an incoming “query” (function f in a Hilbert space), to deliver its best approximator in “value” (trial function space), the “key” (test function space) has to be sufficiently rich such that there exists a key to unlock every possible value.
Dynamic basis update. Another perspective is to interpret the Galerkin-type dot-product attention (14) as a change of basis: essentially, the new set of basis is the column space of Q, and how to linearly combine the bases in Q is based on the inner product (response) of the corresponding feature dimension’s basis in V against every basis in K. From this perspective (Q: values, K: keys, V : query), we have the following result of a layer-wise dynamical change of basis: through testing against the “keys”, a latent representation is sought such that “query” (input trial space) and “values” (output trial space) can achieve the minimum possible difference under a functional norm; for details and the proof please refer to Appendix D.3.4. Theorem 4.4 (layer-wise dynamic basis update, simple version). Under the same assumption as Theorem 4.3, it is further assumed that b(·, q) is bounded below on Kh ⊂ K = V ⊂ H and a(·, ·) : V × K → R is continuous. Then, there exists a set of projection matrices to update the value space {q̃l(·)}dl=1 ⊂ Qh = span{ql(·)}dl=1, for zj ∈ Qh (j = 1, · · · , d) obtained through the basis update rule (14), it holds∥∥a(vj , ·)− b(·, zj)∥∥K′h ≤ minq∈Qh maxk∈Kh |a(vj , k)− b(k, q)|‖k‖K . (16)
The role of feed-forward networks and positional encodings in the dynamic basis update. Due to the presence of the concatenated coordinates x := ‖ni=1xi ∈ Rn×m to the latent representation y, the pointwise subnetwork gs(·) : Rn×m → Rn×d of the nonlinear universal approximator (FFN) in each attention block is one among many magics of the attention mechanism. In every attention layer, the basis functions in Qh/Kh/Vh are being constantly enriched by span{wj ∈ Xh : wj(xi) = (gs(x))ij , 1 ≤ j ≤ d} ⊂ H, thus being dynamically updated to try to capture how an operator of interest responses to the subset of inputs. Despite the fact that the FFNs, when being viewed as a class of functions, bear no linear structure within, the basis functions produced this way act as a building block to characterize a linear space for a learnable projection. This heuristic shows to be effective when the target is assumed to be a function of the (relative) positional encodings (coordinates, transformation groups, etc.), in that this is incorporated in many other attention-based learners with applications in physical sciences [82, 31, 44, 48].

heading: 5 Experiments
text: In this section we perform a numerical study the proposed Fourier Transformer (FT) with the Fourier-type encoder, and the Galerkin Transformer (GT) with the Galerkin-type encoder, in various PDE-related operator learning tasks. The models we compare our newly proposed models with are the operator learners with the simple attention replaced by the standard softmax normalized scaled dot-product attention (ST) [88], and a linear variant (LT) [76] in which two independent softmax normalizations are applied on Q,K separately.2 The data are obtained courtesy of the PDE benchmark under the MIT license.3 For full details of the training/evaluation and model structures please refer to Appendix C.
Instead of the standard Xavier uniform initialization [34], inspired by the interpretations of Theorem 4.3 in Appendix D.3.4, we modify the initialization for the projection matrices slightly as follows
W init ← ηU + δI, for ∈ {Q,K, V }, (17)
where U = (xij) is a random matrix using the Xavier initialization with gain 1 such that xij ∼ U([− √ 3/d, √ 3/d]), and δ is a small positive number. In certain operator learning tasks, we found that this tiny modification boosts the evaluation performance of models by up to 50% (see Appendix C.2) and improves the training stability acting as a cheap remedy to the lack of a softmax normalization. We note that similar tricks have been discovered concurrently in [23].
Unsurprisingly, when compared the memory usage and the speed of the networks (Table 1), the Fourier-type attention features a 40%–50% reduction in memory versus the attention with a softmax normalization. The Galerkin attention-based models have a similar memory profile with the standard linear attention, it offers up to a 120% speed boost over the linear attention in certain tests.
The baseline models for each example are the best operator learner to-date, the state-of-the-art Fourier Neural Operator (FNO) in [57] but without the original built-in batch normalization. All attention-based models match the parameter quota of the baseline, and are trained using the loss in
2https://github.com/lucidrains/linear-attention-transformer 3https://github.com/zongyi-li/fourier_neural_operator
(2) with the same 1cycle scheduler [78] for 100 epochs. For fairness, we have also included the results for the standard softmax normalized models (ST and LT) using the new layer normalization scheme in (5) and (6). We have retrained the baseline with the same 1cycle scheduler using the code provided in [57], and listed the original baseline results using a step scheduler of 500 epochs of training from [57] Example 5.1 and Example 5.2, respectively.

heading: 5.1 Example 1: viscous Burgers’ equation
text: In this example, we consider a benchmark problem of the viscous Burgers’ equation with a periodic boundary condition on Ω := (0, 1) in [57]. The nonlinear operator to be learned is the discrete approximations to the solution operator T : C0p(Ω) ∩ L2(Ω) → C0p(Ω) ∩H1(Ω), u0(·) 7→ u(·, 1). The initial condition u0(·)’s are sampled following a Gaussian Random Field (GRF). The result can be found in Table 2a. All attention-based operator learners achieve a resolutioninvariant performance similar with FNO1d in [57]. The new Galerkin projection-type layer normalization scheme significantly outperforms the regular layer normalization rule in this example, in which both inputs and targets are unnormalized. For full details please refer to Appendix C.2.

heading: 5.2 Example 2: Darcy flow
text: In this example, we consider another well-known benchmark −∇ · (a∇u) = f for u ∈ H10 (Ω) from [10, 57, 56, 64], and the operator to be learned is the approximations to T : L∞(Ω)→ H10 (Ω), a 7→ u, in which a is the coefficient with a random interface geometry, and u is the weak solution. Here L∞(Ω) is a Banach space and cannot be compactly embedded in L2(Ω) (a Hilbert space), we choose to avoid this technicality as the finite dimensional approximation space can be embedded in L2(Ω) given that Ω is compact.
The result can be found in Table 2b. As the input/output are normalized, in contrast to Example 5.1, the Galerkin projection-type layer normalization scheme does not significantly outperform the regular layer normalization rule in this example. The attention-based operator learners achieve on average 30% to 50% better evaluation results than the baseline FNO2d (only on the fine grid) using the same trainer. For full details please refer to Appendix C.3.

heading: 5.3 Example 3: inverse coefficient identification for Darcy flow
text: In this example, we consider an inverse coefficient identification problem based on the same data used in Example 5.2. The input (solution) and the target (coefficient) are reversed from Example 5.2, and the noises are added to the input. The inverse problems in practice are a class of important tasks in many scientific disciplines such as geological sciences and medical imaging but much more difficult due to poor stability [51]. We aim to learn an approximation to an ill-posed operator T : H10 (Ω) → L∞(Ω), u + Nν(u) 7→ a, where Nν(u) stands for noises related to the sampling distribution and the data. = 0.01 means 1% of noise added in both training and evaluation, etc.
The result can be found in Table 3. It is not surprising that FNO2d, an excellent smoother which filters higher modes in the frequency domain, struggles in this example to recover targets consisting of high-frequency traits (irregular interfaces) from low-frequency prevailing data (smooth solution due to ellipticity). We note that, the current state-of-the-art methods [16] for inverse interface coefficient identification need to carry numerous iterations to recover a single instance of a simple coefficient with a regular interface, provided that a satisfactory denoising has done beforehand. The attentionbased operator learner has capacity to unearth structurally how this inverse operator’s responses on a subset, with various benefits articulated in [56, 57, 5, 64, 10].

heading: 6 Conclusion
text: We propose a general operator learner based on a simple attention mechanism. The network is versatile and is able to approximate both the PDE solution operator and the inverse coefficient identification operator. The evaluation accuracy on the benchmark problems surpasses the current best state-ofthe-art operator learner Fourier Neural Operator (FNO) in [57]. However, we acknowledge the limitation of this work: (i) similar to other operator learners, the subspace, on which we aim to learn the operator’s responses, may be infinite dimensional, but the operator must exhibit certain low-dimensional attributes (e.g., smoothing property of the higher frequencies in GRF); (ii) it is not efficient for the attention operator to be applied at the full resolution for a 2D problem, and this limits the approximation to a nonsmooth subset such as functions in L∞; (iii) due to the order of the matrix product, the proposed linear variant of the scaled dot-product attention is non-causal thus can only apply to encoder-only applications.

heading: 7 Broader Impact
text: Our work introduces the state-of-the-art self-attention mechanism the first time to PDE-related operator learning problems. The new interpretations of attentions invite numerical analysts to work on a more complete and delicate approximation theory of the attention mechanism. We have proved the Galerkin-type attention’s approximation capacity in an ideal Hilbertian setting. Numerically, the new attention-based operator learner has capacity to approximate the difficult inverse coefficient identification problem with an extremely noisy measurements, which was not attainable using traditional iterative methods for nonlinear mappings. Thus, our method may pose a huge positive impact in geoscience, medical imaging, etc. Moreover, traditionally the embeddings in Transformerbased NLP models map the words to a high dimensional space, but the topological structure in the same feature dimension between different positions are learned thereby not efficient. Our proof provides a theoretical guide for the search of feature maps that preserve, or even create, structures such as differentiability or physical invariance. Thus, it may contribute to the removal of the softmax nonlinearity to speed up significantly the arduous training or pre-training of larger encoder-only models such as BERT [27], etc. However, we do acknowledge that our research may negatively impact on the effort of building a cleaner future for our planet, as inverse problems are widely studied in reservoir detection, and we have demonstrated that the attention-based operator learner could potentially help to discover new fossil fuel reservoirs due to its capacity to infer the coefficients from noisy measurements.
Acknowledgments and Disclosure of Funding
The hardware to perform this work is kindly donated by Andromeda Saving Fund. The first author was supported in part by the National Science Foundation under grants DMS-1913080 and DMS-2136075. No additional revenues are related to this work. We would like to thank the anonymous reviewers and the area chair for the suggestions on improving this article. We would like to thank Dr. Long Chen (Univ of California Irvine) for the inspiration of and encouragement on the initial conceiving of this paper, as well as numerous constructive advices on revising this paper, not mentioning his persistent dedication of making publicly available tutorials [17] on writing beautiful vectorized code. 4 We would like to thank Dr. Ari Stern (Washington Univ in St. Louis) for the help on the relocation during the COVID-19 pandemic. We would like to thank Dr. Likai Chen (Washington Univ in St. Louis) for the invitation to the Stats and Data Sci seminar at WashU that resulted the reboot of this study. 5 We would like to thank Dr. Ruchi Guo (Univ of California Irvine) and Dr. Yuanzhe Xi (Emory Univ) for the invaluable feedbacks on the choice of the numerical experiments. We would like to thank the Kaggle community, including but not limited to Jean-François Puget (CPMP@Kaggle) for sharing a simple Graph Transformer in TensorFlow,6 Murakami Akira (mrkmakr@Kaggle) for sharing a Graph Transformer with a CNN feature extractor in Tensorflow, 7 and Cher Keng Heng (hengck23@Kaggle) for sharing a Graph Transformer in PyTorch.8 We would like to thank daslab@Stanford, OpenVaccine, and Eterna for hosting the COVID-19 mRNA Vaccine competition and Deng Lab (Univ of Georgia) for collaborating in this competition. We would like to thank CHAMPS (Chemistry and Mathematics in Phase Space) for hosting the J-coupling quantum chemistry competition and Corey Levinson (Eligo Energy, LLC) for collaborating in this competition. We would like to thank Zongyi Li (Caltech) for sharing some early dev code in the updated PyTorch fft interface and the comments on the viscosity of the Burgers’ equation. We would like to thank Ziteng Pang (Univ of Michigan) and Tianyang Lin (Fudan Univ) to update us with various references on Transformers. We would like to thank Joel Schlosser (Facebook) to incorporate our change to the PyTorch transformer module to simplify our testing pipeline. We would be grateful to the PyTorch community for selflessly code sharing, including Phil Wang(lucidrains@github) and Harvard NLP group [52]. We would like to thank the chebfun [28] for integrating powerful tools into a simple interface to solve PDEs. We would like to thank Dr. Yannic Kilcher (ykilcher@twitter) and Dr. Hung-yi Lee (National Taiwan Univ) for frequently covering the newest research on Transformers in video formats. We would also like to thank the Python community [87, 68] for sharing and developing the tools that enabled this work, including PyTorch [69], NumPy [39], SciPy [89], Plotly [45] Seaborn [93], Matplotlib [43], and the Python team for Visual Studio Code. We would like to thank draw.io [46] for providing an easy and powerful interface for producing vector format diagrams. For details please refer to the documents of every function that is not built from the ground up in our open-source software library.9
4https://github.com/lyc102/ifem 5Transformer: A Dissection from an Amateur Applied Mathematician 6https://www.kaggle.com/cpmpml/graph-transfomer 7https://www.kaggle.com/mrkmakr/covid-ae-pretrain-gnn-attn-cnn 8https://www.kaggle.com/c/stanford-covid-vaccine/discussion/183518 9https://github.com/scaomath/galerkin-transformer

Review summaries:
This paper claims to apply two Transformer variants to a data-driven operator learning problem, where the mappings between infinite-dimensional spaces of functions are learned. However, this paper only discusses the uniform grid problem setting, which can be regarded as 0D/1D/2D time series forecasting problem. After applying an interpolation-based CNN network (CiNN) to get $n_c \times n_c \times d$ features, the authors use their proposed Fourier or Galerkin Transformer variants to encode the features and use another CiNN as a decoding regressor. Compared to the previous work FNO, the novelty of the paper comes from using Transformers to encode features instead of applying convolutions by FFT/iFFT. However, the use of Transformers makes the proposed method hard to conduct a zero-shot super-resolution prediction, which can be regarded as the merit of the neural PDE solvers. 

In this paper, they propose a general operator learner based on a simple attention mechanism to remove the softmax operation in standard scaled dot product attention. which is regarded as a PDE-related operator. They also introduce a new scale-preserving layer normalization module to better support their operator learning. Combing with  Fourier Neural Operator, the proposed model obtains great performance on some benchmarks with superior memory and time efficiency.