Choose a Transformer: Fourier or Galerkin
Authors: Shuhao Cao
Venue: NeurIPS 2021 Poster (2021)
Keywords: Transformer
attention
operator learning
partial differential equations
Galerkin methods

Abstract: In this paper, we apply the self-attention from the state-of-the-art Transformer in Attention Is All You Need 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.

Choose a Transformer: Fourier or Galerkin
Shuhao Cao
Department of Mathematics and Statistics Washington University in St. Louis
s. caa@wustl..eu
Abstract
In this paper, we apply the self-attention from the state-of-the-art Transformer Attention Is All You Need [88] for the first time data-driven operator learning problem related to partial differential equations. An effort 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 proved be comparable a Petrov- projection layer-wise, and the estimate is independent with respect to the sequence length. A new layer normalization scheme mimicking the tion is proposed to allow scaling propagate through attention layers, which helps the model achieve remarkable accuracy in operator learning tasks with unnor- malized 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.
Petrov-Galerkin Galerkin
projec-
Introduction
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 study these challenging problems. Traditional methods, such as finite el- ement/difference [20, 22], spectral methods [12], etc., leverage discrete structure to reduce an infinite dimensional operator map 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 Physies-Ifformed Neural Network (PINN) [71, 62, 49] to exploit measurements on collocation points 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. slight change this coefficient invokes potentially expensive re- re-training of any data-driven function leamers. In contrast, an operator learner aims to learn map between finitte-iimessional function spaces, which is much more difficult yet rewarding. A well-trained operator learner can evaluate many instances without re- re-training or collocation points, thus saving valuable resources, and poses itself more efficient approach in the long run. Data-driven resolution-ine -nnvariant operator learning is a booming new research direction (60, 5, 56, 64, 90. 57, 61, 97, 74]. The pioneering model, DeepONet [60], attributes architecturally to universal approximation theorem for operators [18]. Fourier Neural Operator (FNO) [57] notably
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

shows an awing saate-of-the-arr 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 the inputs as targets. Since both functions are sampled discrete grid points, this is special case of a seq2seq problem [81]. The current state-of-the-an seq2seq model the Transformer first introduced 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 better serve the purpose of operator learning.
[88]. As the heart and soul of the Transformer, the scaled
Among our new Hilbert space-inspired adaptations of the scaled dot-product attention, the first and foremost change is: 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 the quadratic complexity of the attention in the classic Transformer. Theory-wise, instead of viewing "row s 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 verify a discrete Ladyzhenskaya-Baauusi Babuška- -abuukka-Brezzi (LBB) condition, which further amounts to the proof that the newly proposed Galerkin-type attention can explicitly represent 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 a cheap diagonal alternative to the normalizations explicitly derived in the proof of the Garee-Gn interpretation (equation (40)). Since a learnable scaling can now be propagated through the encoder layers, the attention- based 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-inspirg projection matrices and layer-wis wise enrichment of the positional encodings, the evaluation accuracies in various operator learning tasks are boosted by significant amount.
tricks including diagonally dominant rescaled initialization for the
Main contributions. The main contributions of this work are summarized as follows.
Attention without softmax. We propose 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 Gatov--alleriin projection.
Operator learner for PDEs. We combine the newly proposed attention operators with the current best sfate-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/memoory efficiency, as well as in accuracy versus those with the conventional softmax normalization. The PyTorch codes to reproduce our results are available as an ooen-source software."
2 Related Works
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-invarian operator learner to achieve a state-of-the-art
http:///hhh....oooccoco///////// tranforrre

---

performance in certain benchmark problems. 190, 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 parametrie 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 additive content-based attention, however, with a vanishing gradient problem due to multiple nonlincarity composition. [25] shows the first effort in removing the softmax normalization in [7] after the projection, however, it still uses 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).
3 Operator learning related to PDEs
Closely following the setup in [56, 57], we consider data-driven model approximate densely- defined operator T : Hi - H2 between two Hilbert spaces with an underlying bounded spacial domain & C 12.h 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 E A, and source f E Yr find u E X such that La Lu) = f.
(i) To approximate the nonlinear mapping from the varying parameter the solution with fixed right-hand side, T : A + X, a + u.
(ii) The inverse coefficient identification problem to recover the coefficient from a noisy measure- ment 7 of the steady-state solution u, in this case, T : X' + A, û + a.
Nonlinear initial value problem: given uo E Ho, find E C((0, T1: H) such that 0.u + N(u) = 0.
(iii) Direct inference from the initial condition to the solution. T Ho + H. vo(-) + u(t1. with h At with i much greater than the step-size in traditional explicit integrator schemes.
Using (i) as an example, based on the given N observations {a6), و and their approximations
{a ()) defined at discrete grid of size h < 1, the goal of our operator learning problem is to =_ =: T(a) s
build an approximation To to T. such that Ta (ax) is a good approximation wa. independent of the mesh size h, where ag and μn are in finite dimensional spaces An, X on this grid. We further assume that (3) ~ for measure compactly supported on A. and the sampled data form reasonably sized subset of A representative of field applications. The loss J I(0) is
[(0] == Enr (IT(a(a)) 1 + 8(a, u: 0)]
(1)
3

---

and in practice is approximated using the sampled observations on discrete grid
J(0)
E {I((T((a))) u/ 1 H 1=1
+ +((a " :
(2)
In example (i): 11 the standard T.2-norm, and ((a, u; 0) serves as regularizer with strength y In Darey flow where La == (aV(.)), it ylla
and is problem-dependeen
since u E H++ (a >> depends on the regularity of a) and a Vv E (div;; St) priori. For the evaluation metric, we drop the ((a, u; f) term, and monitor the minimization of (2) using || 1lx-
Attention-based operator learner
Feature extractor. We assume the functions in both inputs and targets are sampled on uniform grid. In an operator learning problem on Q C IR, simple feedforward neural network (FFN) is used as the feature extractor that is shared by every position (grid point).
Interpolation-based CNN. If C R?, inspired by the multilevel graph kernel network in [56], use two 3-level interpolation-batt based CNNs (CiNN) the feature extractor, but also the downsampling and upsampling layer, respectively, in which we refer to restriction/prolongations between the coarse/fine grids both 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 's input latent representation reside, are concatenated as additional feature dimension(s) to the input, well to each latent representation in every attention head.
Problem-dependent decoder. The decoder is a problem-dependee admissible network that maps the learned representations from the encoder back the target dimension. For smooth and regular solutions in H' (22). we opt for a 2-layer spectral convolution that is the core component in [57]. simple pointwise feedforward neural network (FFN) is used for nonsmooth targets in I (22).
4.1 Simple self-attention encoder
Yis
Q
mmm
softmar
Yoct
yi
Vin
K
1/v2
Attn
V
(a)
A Attng
Y=
K
I.N
add
FFN
add
Yor
Y
1/1r
(b)
Figure 1: Comparison of the vanilla attention [88] with the Galerkin-type simple self-attention in a single head; (a) in the standard softmax attention, the softmax is applied row-wise after the matrix product matmul; (b) a mesh-weighted normalization allows an integration-asee interpretation.
The encoder contains stack of identical simple attention-based encoder layers. For simplicity, we consider single attention head that maps y E ga to another element in gexd and define the

---

trainable projection matrices, and the latent representations Q/K/V as follows.
Wº, w WV E d and Q : :=yye K == :=yww, V=
(3)
We propose the following simple attention that (i) uses a mesh (inverse sequence (ngth)-weightee normalization without softmax, (ii) allows scaling to propagate through the encoder layers.
Ap R - RRxd y y Attn; (y), y - y + +g(),
(4)
where the head-wiss normalizations applied pre-dot-product for t {f.g),
(Fourier-type attention)
Aytn((y) (QK' V/nn
(5)
(Galerkin- type attention) = Att == Q(K V)/n,
(6)
and ō denotes a trainable non-batch-basee normalization. As in the classic Transformer [88], and inspired by the Galerkin projection interpretation, we choose 6 as the layer normalization Ln(o), and 9(-) 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, not all positive anymore. This can be viewed as cheap alternative to the cosine similarity based attention. Remark 4.1. we apply the regular layer normalization rule that eliminates any scaling:
y La(y + Ayin+(y) + o( Ta((y + Attny (x)))) where Aytm; = Q(K (7) then this reduces to the efficient attention first proposed in [76].
4.1.1 Structure-preserviie feature map function of positional encodings
Consider an operator learning problem with underlying domain Q C. R². {-,);=, denotes the set of grid points in the discretized Ω such that the weight 1/n = is the mesh size. Let 1xd denote the feature maps of Q, K, V, i.e., the i-th row of Q, K, V written q: Sg (2x), k; &b (z:), v; = (x;). They are, in the NLP convention, viewed as the feature (embedding) vector at the i-th position, respectively. The inter-position topological structure such continuity diffniiit/iiffeeetttaiitty in the same feature dimension is learned thus not explicit. The following ansatz for Q/ K/V in the same attention head fundamental to our new interpretations. Assumption 4.2. The columns of Q/K Q/K/V, respectively, contain the vector representations of the learned basis functions spanning certain subspaces of the lalent representation Hilbert spaces.
(,(); ou (-), (.)
Using V axd with full column rank as an example, its columns contain potentially set of bases { evaluated at the grid points (degrees of freedom, or DoFs). Similarly, the learned bases whose DoFs form the columns of Q, K are denoted {gr (-)) (ky())) well as {=) for the outputs in (5) and (6). To be specific, the j-th column of V, denoted by wú then stands for a vector representation of the j-th basis function evaluated at each grid point, i.e., its I-th position stands for (+') ) 1y (z)). Consequently, the row v, (1) (z:)..... ua(x;)) can be alternatively viewed as the evaluation of a vector latent basis function at 2z.
4.1.2 Fourier-type attention of quadratic complexity
K
Fil colume O
V
QKT 6-th FOM
FIN COL CAMURA
V
Z
54th
Figure 2: A dissection of Fourier-type attention's 's output. Both matmuls have complexity O(n'd).
In the Fourier-type attention (5), Q, K are assumed to be normalized for simplicity, the j-th column S j S d) in the i-th row (1 < S n)) is computed by (see Figure 2):
(z+); = MQQ ). v¹ = lg(g ki...., ki...., 4: k.
w
n
qi ko) = / (c((x))
(8)
h Σ
1=1
5

---

where the h-weight facilitates the numerical quadrature interpretation of the inner product. Concate- nating columns i S d yields the i-row 2, of the output z: 2; = Jn (L,( OR (E)) (,((E) de. Therefore, without the softmax nonlinearity, the local dot-product attention output at i-th row computes approximately an integral transform with a non-symmetrie learnable kernel function g) == So (5) evaluated at I2, 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 W (1. F2t such that oj + 0 for 1 S j S d, under Assumption 4.2:
(9) This is the forward propagation of the Fredholm equation of the second-kind for each v.)-). When using an explicit orthogonal expansion such as Fourier to solve for (,,(-)) or to seek for a better set of {-)} } in our case, it is long known being equivalent to the Nyström's 's method with numerical integrations [8] (similar to the h 1/nn weighted sum). Therefore, the successes of the random Fourier features in [19, 70] and the Nyströmformer's approximation [97] are not surprising.
07 (x) = 21(2) 2)(2)) <<,,,))s((5 (E)
for 1, dd and E {z.)==r-
Finally, we name this type of simple attention "Fourier" is due to the striking resemblance between the scaled dot-pree dot-product attention and a Fourier-type kernel [30] integral transform, since eventually the target resides in a Hilbert space with underlying spacial domain Ω while the latent representation space parallels a "frequency" domain on .2". This also bridges the structural similarity the scaled dot-product attention with the Fourier Neural Operator [57] where the Fast Fourier Transform (FFT) can be viewed as non-learnable change of basis.
4.1.3 Galerkin-type attention of linear complexity
K
+ith columin
F-th row
Fith column
Q
K
Z
O0
Figure 3: A dissection of Calerkin-type attention's output. Both matmuls have complexity O(nd²).
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 z of (see Figure 3):
(z h h (K
(KTV),,,
(10)
which is the inner product of the i-th row of Q and the j-th column of KVV. Thus,
T
(K V).g = ( (K
q
91
42
(11)
This reads as: (K VV,, contains the coefficients for the linear combination of the vector represen- tations {g' [==1 of the bases stored in Q's column space to form the output z. Mcanwhile, the j-th column (K V).j of K TV consists the inner product of j-th column of V with every column of K.
1d
d
21 h Σ o (KTTV) where (K (KTV).; (k= v. 122 v.
ka. v') k.v''
(12)
As 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., kr(:), ty(+) evaluated at every 2; are simply their vector representations k (1 < l < d) and v', we have the functions represented by the columns of the output z can then compactly written as: rewriting (1y, kt) == (KVV,,
d
(z) = Σ
(uj, ki) q((x), for j = 1,... .....,... ,d, and æ E (1:)
(13)

---

where the bilinear form (;,)) x K - R. (13) can be also written in a componentwise form:
d
(%s (s(()Ri(() de (xx)).
(z²); 25(x;) = ); - h Σ (k' a (d)
= 1=1
Σ
(14)
1=1
Therefore, when (o) E (q. k u) 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 that of K, respectively, the forward propagation of the Galerkin-type attention is a recast of a learnable Petrov-Galerkin-yype projection (cf. Appendix D.I) 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 Q C I.th discretized by n grid points, and J E H. y E Raxd is the current latent representation for n > d > in and full column rank. Qn Q C H and Vn. C VV C H are the latent approximation subspaces spanned by busis functions with the columns of Q and V in (3) as
degrees of freedom, respectively, and O < dim Qh = r S dim Vn = d. Let o(.-)): V Q Q-2 R be continuous bilinear form, and if for any fixed E Qn the functional norm of b(: q) is bounded below by с > O, then there exists a learnable map go that is the composition of the Galerkin-type
attention operator with updated set of projection matrices (We, W WV), and a pointwise
universal approximator, such that for fr. E Qn being the best approximation of f in Il it holds:
bb(, 9)1 + llf - -fallxe
|l - -el Sc min max qEE, DEV,
(15)
|00|
Remarks 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 || llx in the current subspace Qn. 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 E 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 e is independent of n in the lower bound of [bb(, I| The existence of an n- r-inddppnndent Baw.kaaaaaaaaaaaaaaaaaaa (LBB) condition [21, Chapter 6. 6.1,], also referred as the Banach-Nečas Nečas-Babbuk (BNB) condition in Galerkin methods on Banach spaces [29, Theorem 2.6].
lower bound is commonly known as the discrete version of the
As the cornerstone of the approximation to many PDEs, the discrete LBB condition establishes the surjectivity of a map from V, to Qa. In a simplified context (15) above of approximating functions using this linear attention variant (Q: values, query, V: keys), roughly translates to: for an incoming "query" (function f in 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 based on the inner product (response) of the corresponding feature dimension's '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 K, C K == C H and a(-s-): ; V x IC- + R is continuous. Then, there exists set of projection matrices to update the value space { C Qn = =span (q(()) for 23 E Qu (j 1, ,d) obtained through the basis
update rule (14), it holds
Ja(v;, k) b(k,g)) S min max qeQu, kERn
||a(,);)) 6(, ((,,,,)0ls
(16)
|
7

---

The role of feed-forward networks and positional encodings in the dynamic basis update. Due to the presence of the concatenated coordinates x == 1";=1=; E R=** the latent representation y, the pointwise subnetwork ga (-): -=xm - Rxdd 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 / / are being constantly enriched by span w E Xx w)(2;) (g. (x)),, S j I d) C 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 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].
5 Experiments
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 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. 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 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 + nU + 81, for o E (Q.K
(17)
where U (1zy) is random matrix using the Xavier initialization with gain such that Tay ~ u([-3/d, /3/d), and 5 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 softmax normalization. The Galerkin attention- based models have a similar memory profile with the standard linear attention, it offers up 120% speed boost over the linear attention in certain tests.
Table 1: The memory usage/FL.OPPcoomlesii comparison of the models. Batch size: 4; the CUDA mem (GB): the sum of the self cuda_ www.memorry... _usage; GFLOP: Giga FLOP for backpropagation (BP); both are from the Py PyTorch autograd profiler for BP averaging from 1000 BPs; the mem (GB) is recorded from nvidia smi of the memory allocated for the active Python process during profiling; the speed (iteration per second) is measured during training; the exponential operation assumed to have an explicit complexity of c. > 1 [14].
cuda_memoryy__....
Example 1: n
8192
Encoders only: a
8192, d 128, I
10
Computational complexity of the dot-product per layer
Mem
CUDA Mem
Speed
GFLOPP
Mem
CUDA Mem
Speed
GELOP
ST
18.39
31.06
5.02
1393
18.53
31.34
4.12
1876
O(n2a,d) O(n'd) O'n(1 - cadd) O(nd)
FT
10.05
22.92
6.10
1138
10,80
22.32
5.46
1610
LT
2.55
2.31
12.70
606
2.73
2.66
10.98
773
GT
2.36
1.93
27.15
275
2.53
2.33
19.20
412
The baseline models for each example are the best operator learner to-date, the state-of-the-ae 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
htt.pp////ttthaaa////cceeeeeaeee transformer
httpp://ttaabboooooogggg////o/oo http://ouurrrrrrruuaaaoeoppepe operator
8

---

(2) with the same 1cycle scheduler [78] for 100 epochs. For faimess, 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 [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.
5.1 Example 1: viscous Burgers' equation
In this example, we consider a benchmark problem of the viscous Burgers' equation with periodic
boundary condition on $ = (0, 1) in [57]. The nonlinear operator be learned is the discrete approximations to the solution operator T' (22) n (2) - C2(2) n H (22, wo(-) ~ u(s 1).
The initial condition no (-)'s I's are sampled following a Gaussian Random Field (GRF).
The result can be found in Table 2a. All attention-based operator learners achieve resolution- invariant performance similar with FNOId [57]. The new Galerkin projection-type layer normal- ization 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.
5.2 Example 2: Darcy flow
In this example, consider another well-known benchmark - (aVu) = = (S) from [10, 57, 56, 64], and the operator to be learned is the approximations to T" (9) + Ho (?t), + a. which a the coefficient with a random interface geometry, and
the weak solution. Here L (52) is Banach space and cannot be compactly embedded in (?( (a Hilbert space), we choose avoid this technicality the finite dimensional approximation space can embedded in L' given that Ω is compact.
The result be found in Table 2b. As the input/output normalized, in contrast Example 5.1, the Galerkin projection-type layer normalization scheme does not significantly outperform the regular layer normalization rule in this example. The attention- basedd 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.
Table 2: (a) Evaluation relative error (x10-3) of Burgers' equation 5.1. (b) Evaluation relative error (x (x100=) of Darcy interface problem 5.2.
(a)
(b)
72
512
2048
n 8192
Rg = 141, 43
ny. n1, 211, 61
FNOId [57]
15.8
14.6
13.9
FNO2d [57]
1.09
1.09
FNOId 1cycle
4.373
4.126
4.151
FNO2d 1cycle
1.419
1.424
FT regular La
1.400
1.477
1.172
FT regular Ln
0.838
0.847
GT regular Ln
2.181
1.512
2.747
GT regular Ln
0.894
0.856
ST regular La
1.927
2.307
1.981
ST regular I.
1.075
1.131
LT regular Ln
1.813
1.770
1.617
LT regular Ln
1.024
1.130
FT Ln on Q, K
1.135
1.123
1.071
FT La on Q,
0.873
0.921
GT Ia on K
1.203
1.150
1.025
GT La K, V
0.839
0.844
ST Ln on Q,
1.271
1.266
1.330
on Q,
0.946
0.959
Ln on K.
1.139
1.149
1.221
LT on
0.875
0.970
5.3 Example 3: inverse coefficient identification for Darcy flow
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 class of important tasks 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
an ill-posed operator T: (2)) - T.x( u (N) + a, where (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 this example to recover targets consisting of traits (irregular interfaces) from low-frequency prevailing data (smooth solution due to ellipticity). We note that, the current state-of-the-ar methods [16] for inverse interface coefficient identification need to carry numerous iterations to recover single instance of simple coefficient with a regular interface, provided that a satisfactory denoising has done beforchand. The attention- based operator learner has capacity to unearth structurally how this inverse operator's responses subset, with various benefits articulated in [56, 57, 5, 64, 10].
Table Evaluation relative error (x10- 3) of the inverse problem 5.3.
Rf. 36
By.
211, 71
0.01
c 0.1
==
0.01
0.1
FNO2d (only ny)
13.71
13.78
15.12
13.93
13.96
15.04
FNO2d (only n2))
14.17
14.31
17.30
13.60
13.69
16.04
FT regular Ln
1.799
2.467
6,814
1.563
2.704
8.110
GT regular Ln
2.026
2.536
6.659
1.732
2.775
8.024
ST regular Ln 2.434
3.106
7.431
2.069
3.365
8.918
LT regular In 2.254
3.194
9.056
2.063
3.544
9.874
FT Ln on Q. K 1.921
2.717
6.725
1.523
2.691
8.286
GT K, V 1.944
2.552
6.689
1.651
2.729
7.903
Ln Q. K 2.160
2.807
6.995
1.889
3.123
8.788
LT La K,
2.360
3.196
8.656
2.136
3.539
9.622
6 Conclusion
We propose general operator learner based on 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-of- the-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 -owdiimesiional attributes (e.g., smoothing property of the higher frequencies in GRF); (ii) it is not efficient for the attention operator to applied at the full resolution for 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- dot-product attention is non-causal thus can only apply to encoder-only applications.
7 Broader Impact
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 more complete and delicate approximation theory of the attention mechanism. We have proved the Galerkin-type attention's 's approximation capacity in an ideal Hilbertian setting. Numerically, the new attention-based operator learner has capacity 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 huge positive impact in geoscience, medical imaging, etc. Moreover, traditionally the emheddings in Transformer- based 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 theoretical guide for the search of feature maps that preserve, or even create, structures such as differentiability physical invariance. Thus, it may contribute to the removal of the softmax nonlinearity speed up significantly the arduous training or pre-traing 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 building 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 capacity to infer the coefficients from noisy measurements.
10

---

Acknowledgments and Disclosure of Funding
The hardware to perform this work kindly donated by Andromeda Saving Fund. The first author was supported in part by the National Science Foundation under grants No additional revenues are related 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 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. We would like to thank Dr. Ari Ster (Washington Univ in St. Louis) for the help on the relocation during the COVID-19 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. We would like to thank Dr. Ruchi Guo (Univ of California Irvine) and Dr. Yuanzhe Xi (Emory Univ) for the invaluable feedbacks the choice of the numerical experiments. We would like to thank the Kaggle community, including but not limited to Jean- -aan-François Puget (CPMP@ Kaggle) for sharing a simple Graph Transformer in TensorFlow, o Murakami Akira (mrkmakr@ Kaggle) for sharing a Graph Transformer with a CNN feature extractor in Tensorflow, and Cher Keng Heng (hengek23 @ Kaggle) for sharing Graph Transformer in PyTorch. We would like thank daslab @ Stanford, OpenVaccine, Vaccine, and Eterna for hosting the COVID-19 mRNA Vaccine competition and Deng Lab (Univ of Georgia) for collaborating this competition. We would like to thank CHAMPS (Chemistry and Mathematies 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 Py PyTocc community for selflessly code sharing, including Phil Wang(lucidrains @ github) and Harvard NIPP group [52]. We would like to thank the chebfun [28] for integrating powerful tools into 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 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- pource software library.
and DMS 21 21365
'nttpe://github.com com/1yc102/ifem
http://ggtaab.coom/yy10////////
'Transformer: A Dissection from Amateur Applied Mathematician
"attps://ww..kag""" Saa.ggll..co/mppppppgggaaaaatata
http://wwwwwkwgll.e
'attps://githun http://cooottt//ggaaaaaaaaaaeaea hwa.ttttthaatteeeeeettateertreer
11

---

References
[1] Jonas Adler and Ozan Öktem. Solving ill-posed inverse problems using iterative deep neural networks. Inverse Problems, 33(12):124007, 2017.
[2] Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character- level language modeling with deeper self-attention. Proceedings of the AAAI Conference on Artificial Intelligence, 33 (01):3159-3144) Jul. 2019.
[3] Giovanni Alessandrini. An identification problem for an elliptic equation in two variables. Annali di matematica pura ed applicata, (45(1))265-295 1986.
[4] Ferran Alet, Adarsh Keshav Jeewajee, Maria Bauza Villalonga, Alberto Rodriguez, Tomas Oozan-PPerez, and Leslie Kaelbling. Graph element networks: adaptive, structured computation and memory. In International Conference on Machine Learning, pages 212 2222222.222 PMLR, 2019.
[5] Anima Anandkumar, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Nikola Kovachki, Zongyi Li. Burigede Liu, and Andrew Stuart. Neural operator: Graph kernel network for partial differential equations. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2020.
[6] Daniel Amdt, Wolfgang Bangerth, Bruno Blais, Thomas C. Clevenger, Mare Fehling, Alexander V. Grayver, Timo Heister, Luca Heltai, Martin Kronbichler, Matthias Maier, Peter Munch, Jean-Paul Pelteret, Reza Rastak, Ignacio Thomas, Bruno Tureksin, Zhuoran Wang, and David Wells. The deal. II library, version 9.2. Journal of Numerical Mathematics, 28(3):131-146, 2020.
[7] Dzmitry Bahdanau, Kyung Hyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, 2015.
[8] lean-Paul Berrut and Manfred R Trummer. Equivalence of Nyström's method and Fourier methods for the numerical solution of Fredholm integral equations. Mathematics of computation, 48(178):617-62)
1987.
[9] Saakaar Bhatnagar, Yaser Afshar, Shaowu Pan, Karthik Duraisamy, and Shailendra Kaushik. Prediction of aerodynamic flow fields using convolutional neural networks. Computational Mechanics, 64(2):525-545, 2019.
[10] Kaushik Bhattacharya, Bamdad Hosseini, Nikola Kovachki, and Andrew M Stuart. Model reduction and neural networks for parametric PDEs. arXiv preprint trXiv:2005.031880 2020.
[11] Guy Blanc and Steffen Rendle. Adaptive sampled softmax with kernel based sampling. In International Conference on Machine Learning, pages
PMLR, 2018
[12] Ronald Newhold Bracewell and Ronald N Bracewell. The Fourier transform and its applications, volume 31999. McGraw-Hie Hill New York, 1986.
[13] Susanne C Brenner and Ridgway Scott. The mathematical theory of finite element methods, volume 15. Springer, 2008
[14] Richard Brent. Multiple- zero-finding methods and the complexity of clementary function evaluation, In Analytic computational complexity, pages 151-176. Elsevier, 1976.
[15] Shuhao Cao, Long Chen, and Ruchi Guo. virtual finite element method for two dimensional Maxwell interface problems with a background unfitted mesh. Mathematical Models and Methods in Applied Sciences, to appear, 2021.
[16] Tony F Chan and Xuc-Cheng Tai. Identification of discontinuous coefficients in elliptic problems using total variation regularization. SIAM Journal Scientific Computing, 25(3):881-908 2003.
[17] Long Chen. 4FEM: an integrated finite element methods package in MATLAB. Technical report, 2008
[18] Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions on Neural Networks, 6(4):911-91 1-917, 1995.
[19] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with Performers. In International Conference Learning Representations (ICLR), 2021.
12

---

[20] Philippe G Ciarlet. The finite element method for elliptic problems. SIAM, 2002.
[21] Philippe G Ciarlet. Linear and nonlinear functional analysis with applications, volume 130. SIAM, 2013.
[22] Richard Courant, Kurt Friedrichs, and Hans Lewy. On the partial difference equations of mathematical physics. IBM journal of Research and Development, 1 11(2)215-234, 1967.
(23) Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuher. The devil in the detail: Simple tricks improve systematic generalization of transformers. In Proc. Conf. on Empirical Methods in Natural I Language Processing (EMNIP), Punta Cana, Dominican Republic, November 2021.
[24] Niccolò Dal Santo, Simone Deparis, and Luca Pegolotti. Data driven approximation of parametrized pdes by reduced basis and neural networks. Journal of Computational Physics, 416: 109550, 2020.
[25] Alexandre de Brébisson and Pascal Vincent. A cheap linear attention mechanism with fast lookups and fixed-size representations. arXiv preprint arXiv: (609000866, 2016.
[26] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019.
[27] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL- 2019, Minneapolis, MN, USA, June 2-7, 2019, Volune 1, pages 4171 4186. Association for Computational Linguistics, 2019.
[28] Tobin A Driscoll, Nicholas Hale, and Lloyd N Trefethen. Chebfun guide, 2014.
[29] Alexandre Ern and Jean-Lue Guermond. Theory and Practice of Finite Elements. Springer, 2004.
[30] Charles Fox. The G and H functions as symmetrical Fourier kernels. Transactions of the American Muthematical Society, 98(3):395 429.99999999110 1961.
[31] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Srany-rransommer:: 3D Roto- Translation Equivariant Attention Networks. In Advances in Neural Information Processing Systems, volume 33, pages 1970-1981, 2020.
[32] D. Gilbarg and N.S. Trudinger. Elliptic Partial Differential Equations of Second Order. Classics in Mathematics. Springer Berlin Heidelberg, 2001. ISBN 7783540411604
[33] Vivette Girault and P-A Raviart. Finite element approximation of the Navier-Stokes equations. Lecture Notes in Mathematics, Berlin Springer Verlag, 749. 1979.
[34] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedingy of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 Proceedings of Machine Learning Research, pages 249-256, Chia Laguna Resort, Sardinia, Italy, 13-15 15 May 2010. PMLR.
[35] Ruchi Guo and Jiahua Jiang. Construct deep neural networks based on direct sampling methods for solving electrical impedance tomography. SIAM Journal on Scientific Computing, B373),B67 2021.
[36] Xiaoxiao Guo, Wei Li. and Francesco lorio. Convolutional neural networks for steady flow approximation. In Proceedings of the 22nd ACM SIGKDD international conference knowledge discovery and data mining, pages 481-490. 2016.
[37] Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet- baasd operator learning for differential equations. arXiv preprint 2Xiv:2109.11454 2021.
[38] W. Hackbusch. Multi-Grid Methods and Applications. Springer Series in Computational Mathematics. Springer Berlin Heidelberg, 2013. ISBN 9783662024270.
[39] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser. Julian Taylor. Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río. Mark Wiebe, Pearu Peterson, Pierre Géraart-Marhhan,, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer A Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature 585(7825):357-36) September 2020.
[40] Juncai He and Jinchao Xu. Mgnet: A unified framework of multigrid and convolutional neural network. Science China Mathematics, 12(7))1331111333 1354, 2019.
13

---

[41] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 630-645. Springer. 2016.
[42] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359-366, 1989.
[43] J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, 93):900955 2007.
[44] Michael J Hutchinson, Charline Le Lan, Shcheryar Zaidi, Emilien Dupont, Yee Whye Tch, and Hyunjik Kim. LieTrrasformrr: equivariant self- for Lic groups. pages 4533-4543, 2021.
[45]
Plotly Technologies Inc. plotly, 2015.
[46] JGraph. draw.io, 2021.
[47] Jiahua Jiang, Yi Li, and Ruchi Guo. Learn an index operator by enn for solving diffusive optical tomography: deep direct sampling method. arXiv preprint arXiv:2/04.07703 2021.
[48] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with AlphaFold. Nature, 59677873)) 2021.
[49] George Em Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 2021.
[50] Angelos Katharopoulos, Apoory Vyas, Nikolaos Puppas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning. pages 5156-5165, PMLR, 2020.
[51] A. Kirsch. An Introduction to the Mathematical Theory of Inverse Problems. Applied Mathematical Sciences. Springer New York, 2011. ISBN 9781441984766.
[52] Guillaume Klein, Yoon Kim, Yuntian Deng, Jean Senellart, and Alexander M. Rush. OpenNMT: Open- Source Toolkit for Neural Machine Translation. In Proc. ACI, 2017.
[53] Peter D Lax and Robert D Richtmyer. Survey of the stability of linear finite difference equations. Communications on pure and applied mathematics, 9(2):267-293 1956.
[54] James Lee- Tho-Th Joshua Ainslie, Ilya Fckstein, and Santiago Ontanon. FNet: Mixing tokens with Fourier transforms. arXiv preprint arXiv:2105.00322 2021.
[55] Sijing Li, Zhiwen Zhang, and Hongkai Zhao. A data-driven approach for multiscale elliptic PDEs with random coefficients based on intrinsic dimension reduction. Multiscale Modeling & &Simulation, 18(3): 1242-1271, 2020.
[56] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. In Advances in Neural Information Processing Systems, volume 33, pages
2020.
[57] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Arizzadencsheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021.
[58] Hongzhou Lin and Stefanie Jegelka. Resnet with one-neuron hidden layers is a universal approximator. Advances in Neural Information Processing Systems (NIPS 2018), 31:6169-6177 2018.
[59] Kevin Lu, Aditya Grover, Pieter Abbeel, and Igor Mordatch. Pretrained transformers as universal computation engines. arXiv preprint arXiv:2103.0522 103.05247, 2021.
[60] Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identi- fying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910.03193 2019.
[61] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3::118-229 2021.
[62] Lu Lu. Xuhui Meng. Zhiping Mao, and George Em Karniadakis. Deepxde: A deep learning library for solving differential equations. SIAM Review, 63(1):208 -28, 2021.
14

---

[63] Peter Monk et al. Finite element methods for Maxwell's equations. Oxford University Press, 2003.
[64] Nicholas H Nelsen and Andrew M Stuart. The random feature model for input-output maps between banach spaces. SIAM Journal on Scientific Computing. 43(5):A3212 A3243, 2021.
[65] Tan M. Nguyen, Vai Suliafu, Stanley J. Osher, Long Chen, and Bao Wang. FMMformer: Efficient and Flexible Transformer via Decomposed Near-field and Far-field Attention. In Advances Neural Information Processing Systems (NeurIPS), 2021.
[66] Toan Q Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self- attention. arXiv preprint arXiv:1910.058%. 2019.
[67] Akifumi Okuno, Tetsuya Hada, and Hidetoshi Shimodaira. A probabilisttic framework for multi-view feature learning with many-to-many associations via neural networks. In International Conference on Machine Learning, pages 3888-3897. PMLR, 2018.
[68] Travis E. Oliphant. Python for scientific computing, Computing in Science Engineering, 9(3):10-20 2007.
[69] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), pages 8024-8035, 2019.
[70] Hao Peng, Nikolaos Pappas. Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference Learning Reprexentations, 2021.
[71] M. Raissi, Perdikaris. and G.E. Karniadakis. Physics-informer neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
[72] Prajit Ramachandran, Barret Zoph, and Quoe Le. Searching for activation functions. arXiv preprint 1rXiv:1710.05941, 2017.
[73] Ankit Singh Rawat, Jiecao Chen, Felix Xinnan X Yu, Ananda Theertha Suresh, and Sanjiv Kumar. Sampled softmax with random fourier features. In Advances in Neural Information Processing Systems, volume 32, 2019.
[74] Nicholas Roberts, Mikhail Khodak, Tri Dao, Liam Li, Christopher Ré, and Ameet Talwalkar. Rethinking neural operations for diverse tasks, arXiv preprint arXiv:2103.15798 2021.
[75] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers, 2021.
[76] Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 3531-3539, 2021.
[77] Jonathan W Siegel and Jinchao Xu. Approximation rates for neural networks with general activation functions. Neural Networks. 128:313-321 2020.
[78] Leslic N Smith and Nicholay Topin. Super- convergeee: Very fast training of neural networks using large learning rates. In Artificial Intelligence and Machine Learning, for Mulri-Domain Operations Applications, volume 1006, page 1006612. International Saciety for Optics and Photonies, 2019.
[79] Kyungwoo Song, Yohan Jung, Dongjun Kim, and 11-Chul Moon. Implicit kernel attention. Proceedings of the AAAI Conference on Artificial Intelligence, 35( 35(11):91339772) May 2021.
[80] Ari Stern, Banach space projections and Gaaoo-G estimates. Numerische Mathematik, 130(1): 125-133, 2015.
[81] Ilya Sutskever, Oriol Vinyals, and Quoe Le. Sequence to sequence learning with neural networks. In Proceedings of the 27th International Conference on Neural Information Processing Syxtems - Volume 2, NIPS'14, 14, page 3 3104-3112, Cambridge, MA, USA, 2014. MIT Press.
[82] Kai Sheng Tai, Peter Bailis, and Gregory Valiant. Equivariant transformer networks. In International Conference Machine Learning, pages 6086 6095. PMLR, 2019.
15

---

[83] Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. MLP-mixer: An all-MLP architecture for vision. arXiv preprint arXiv:2105.0100 2021.
[84] Yao- Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philiipe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unitied understanding for transformer's s attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIP) pages 4344-4353, Hong Kong, China, November 2019.
[85] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv: 1607. 08022. 2016.
[86] Benjamin Ummenhofer, Lukas Prantl, Nils Thuerey, and Vladlen Koltun. Lagrangian fluid simulation with continuous convolutions. In International Conference Learning Representations, 2020.
[87] Guido Van Rossum and Fred L. Drake. Python 3 Reference Manual. CreateSpace. Scotts Valley, CA. 2009. ISBN 4414126677.
[88] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NIPS 2017), volume 30, 2017.
[89]
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapcau, Evgeni Burovski, Pearu Peterson, Warren Weekesser, Jonathan Bright, Stéfan van der Walt, Matthow Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, Ilhan Polat, Yu Feng. Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SeiPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020.
[90] Sifan Wang. Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal on Scientific Computing. 4365):3305--3381,, 2021.
[91] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learinng the solution operator of parametric partial differential equations with physics-inforee informed DeepOnets. arXiv preprint arXiv:2103. 2021.
[92] Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 04768, 2020.
[93] Michael I.. Waskom. scaborn: statistical data visualization. Journal of Open Source Software, 6(60); 3021, 2021.
[94] H. Whitney. Geometric Integration Theory. Princeton Legacy Library. Princeton University Press, 2015. ISBN 9781400877577.
[95] Matthew Wright and Joseph E Gonzalez. Transformers are deep infnittediimesiional non- non-Mereer binary kemel machines. arXiv preprint (xXi::2006.01506 2021.
[96] Ruihin Xiong, Yunchang Yang, Di Hc, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang. and Ticyan Liu. On layer normalization in the transformer architecture. International Conference on Machine Learning, pages 10524-10533. PMI.R, 2020.
[97] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A Nyström-based algorithm for approximating self-attention. Proceedings of the AAAI Conference Artificial Intelligence, 35(16): 14138-1488, May 2021.
[98] Jinchao Xu and Ludmil Zikatamov. Algebraic multigrid methods. Acta Numerica, 26:591-721, 2017.
[99] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeence functions? In International Conference on Learning Representations, 2020.
[100] Yinhao Zhu and Nicholas Zabaras, Bayesian deep convolutional encoder networks for surrogate modeling and uncertainty quantification. Journal of Computational Physics, 366:415-447, 2018
16

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In Section 6 and in each
experiment's ' detailed description in Appendix C.3 and C.4.
(c) Did you discuss any potential negative socictal impacts of your work? [Yes] In Section
7.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] The key assumption of the new interpretations of the attention mechanism is in Assumption 4.2. Due to the page limit, the assumptions and settings are fully elaborated in mathematical rigorous fashion in Appendix for the proof of Theorem 4.3.
(b) Did you include complete proofs of all theoretical results? [Yes] In Appendix D.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material The data are publicly available. For instruction to reproduce our result please re- fer to README. md in https:///gtthti comb coom//scomaah////oooooooooaaaaaa wetpt//ggllekkinaaeeee trankin--ttanffoomrr... tree/maing /xamm//xxmpss....
URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] We include all details in the supplemental material Appendix B and C.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] We found that our method is numerically seed-invariant, and we have reported the error bands in Appendix C.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Please refer Appendix C.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yss] Multiple creators of or contributing work to the dataset used in our work are cited [57. 5, 56, 64]. The code base we have implemented our model upon is the footnote of Section and cited in the Acknowledgement [52]. In the publicly available code repository, in the document of every function, we have credited every author we can find for even small code snippet, if that function is not built from ground up by us.
(b) Did you mention the license of the assets? [Yes] In Section 5.
(c) Did you include any new assets either in the supplemental material The full code base to replicate our results publicly available as an open-source software.
URL? [Yes]
(d) Did you discuss whether and how consent obtained from people whose data you're using/curating? [Yes] The data obtained from https : //github. com com/zongyi-li. f fourier_neural_ooo wet.inneelleerrttttteeee author per a personal communication.
with the MIT license well as the consent from the
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
17


Reviews:
[{"review_id": "wFPd2dTHUAf", "reviewer": "Reviewer_xknS", "summary": "This paper aims to improve the well-established attention mechanism, and challenges the necessity of its softmax normalization through an integral perspective.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper aims to improve the well-established attention mechanism, and challenges the necessity of its softmax normalization through an integral perspective. The outcome are two “softmax-free” attention variants, connected respectively to Fourier feature transform (and thereby to several recent work sharing this insight, as noted by the paper) and Petroc-Galerkin projection. For the latter, non-trivial theoretical results on approximation error are presented. The proposed attention variants have practical benefits too: in several operator learning experiments, they outperform baselines in terms of accuracy and efficiency.\n\nOverall I appreciate the paper’s technical significance, and really like it that the paper connects to previous works every now and then. Although the presentation is reasonably clear, I do have some confusion/suggestions: \n- Replacing the softmax with a linear-ish interaction hints significant efficiency gain. The proposed method achieves speedup and memory saving over baselines, and the paper would appeal to a broader audience if it promises so upfront. \n- I would appreciate it if the narrative does a better job connecting to what the implantation looks like. E.g., walking the readers through a quick algorithm would do. \n- Thm. 3.3 proceeds with a discrete space with length n, but isn’t bilinear form is over the Hilbert spaces?  I am not an expert in this direction and did not have the capacity to checkout out the proof, so I might be missing something pretty obvious. It would be great if the authors can clarify.\n- Adding onto the above, Thm. 3.3 definitely stills has its value, since in practice the spaces are discretized anyway. But it is intriguing to see that the bound actually does not depend on n. I’m curious whether or not there is a more general version of 3.3, that could hold for the infinite-dimensional Hilbert-space case?\n\nOverall the paper is solid and timely. I vote for acceptance.\n\nStrengths:\n- Interesting and timely contribution to “linearizing” attention in transformers.\n- Well-designed experiments and solid results.\n- The theoretical treatment of attention could inspire future research.\n\nWeaknesses:\n- The paper could do a better job in connecting to implementation and benefits in practice\n- The key theoretical results (Thm. 3.3) needs further clarification.\n\nMissing reference:\nhttps://arxiv.org/abs/2103.02143", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}, {"review_id": "QLdXoxLW3fo", "reviewer": "Reviewer_3j15", "summary": "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.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nAdvantages:\n1. The interpretations of non-softmax attention as Fourier expansion and Petrov–Galerkin-type projection seem to be interesting and insightful, while the proposed methods are not novel compared to other linear attention variants.\n2. The proposed Fourier or Galerkin Transformer variants show some improvements over the standard softmax normalized scaled dot-product attention and the linear variant on operator learning problems.\n\nWeaknesses:\n1. While the application of random Fourier features on Transformers is discussed in the paper, neither Performer nor RFA is compared in the experiments.\n2. Unlike the previous work FNO, the use of Transformers makes the proposed method hard to conduct zero-shot super-resolution prediction. In this case, whether the proposed method can be regarded as a neural PDE solver is questionable. And the paper doesn't present any experiments w.r.t this point (i.e., zero-shot super-resolution prediction).\n\nPossible Typo:\n\nline 105: $y \\rightarrow y + \\sigma(y + \\mathrm{Attn}(y))$ are not consistnet with the original definition of Transformer. It could be $y \\rightarrow y + \\mathrm{Attn}(y) + \\sigma(y + \\mathrm{Attn}(y))$ if possible.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "8WjhaiiZXv5", "reviewer": "Reviewer_48dv", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is written clearly and they regard the standard attention mechanism in Transformer as an operator learning problem with removing the softmax operations. Extensive experiments are conducted to confirm the effectiveness and efficiency of their proposed method.   It would be interesting if this proposed method could be integrated into pre-train language models for speeding up.  Just have a few questions:\n1. Could the proposed method be applied to tasks like machine translation or GLUE?\n2. Is the proposed method able to model the autoregressive attention such as in NMT or language modelling?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, they adequately address the limitations.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "-MA_q6M1SZ", "reviewer": "Reviewer_wufR", "summary": "This paper proposes two Transformer variants for data-driven operator learning problems of parametric partial differential equations: the Fourier-type and the Galerkin-type. This paper includes both theoretical evidence and empirical evidence to show that these two Transformer variants are both accurate and efficient for solving various parametric PDEs.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper first generalizes the Transformer architecture and then shows that vanilla Transformer, Fourier-type Transformer, Galerkin-type Transformer are special cases of this generalized Transformer architecture. This generalization is natural but creative. The author uses informative diagrams and equations to illustrate different kinds of Transformers talked about in this paper clearly.\n\nThe evaluation is on three well-known PDE benchmarks. This paper compares the error of all different combinations of Layer Normalization, so the experiment section is comprehensive and solid.\n\nAlthough the variants seem to be inspired by the mathematical representations of some integrations, the proposed variants have only a few differences from the standard Transformer in the source code. I suggest that the author can list all these differences explicitly to make them clear. To my understanding, the differences only lie in three places:\n\n1. Use mesh-weighted normalization without softmax\n2. Where to add layer normalizations are different\n3. (For Galerkin-type Transformer only) the ordering of matrix multiplication is different\n\nThe author does not explain why the speed/memory of ST, FT, LT, and GT are different in Table 1. To my understanding, LT and GT are faster than ST and FT because they are linear, but why is FT faster than ST, GT faster than LT? It would be better if there are more detailed explanations (maybe manually calculated FLOPs/memory).\n\nTo summarize, the motivation of this paper is interesting, but the real difference of the proposed methods with standard Transformer is limited and is not shown clearly. Shown by solid experimental results, the proposed variants are both accurate and efficient in solving PDEs, but the efficiency difference requires more explanations.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies Transformer variants for data-driven operator learning problems of parametric partial differential equations: the Fourier-type and the Galerkin-type. This paper includes both theoretical and empirical evidence to show that these two Transformer variants are accurate and efficient for solving various parametric PDEs.

The authors have nice discussions with all the reviewers during the rebuttal period and put great efforts into improving the paper quality. All the reviewers agree that the mapping between Transformer and PDE solvers is interesting and important. Therefore, I recommend acceptance. I hope the authors can revise the paper in the camera-ready version.
Decision: Accept (Poster)