Choose a Transformer: Fourier or Galerkin

1 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 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 can not 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 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-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 ] . 3 Operator learning related to PDEs . 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 .

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.