PaperlayerView exact revision on arXiv

Source-linked AI summary

Choose a Transformer: Fourier or Galerkin

Shuhao Cao

arXiv:2105.14995v4cs.LGmath.NA

TL;DR

The paper addresses efficient learning of PDE solution and inverse operators, which are harder than single-instance function approximation and especially sensitive to measurement noise. It introduces softmax-free attention with a Galerkin-inspired normalization and reports improved evaluation accuracy over FNO and softmax-normalized counterparts, including noisy inverse identification.

  • Problem

    Operator learning seeks reusable maps between infinite-dimensional function spaces, while inverse operators are less stable and amplify measurement noise.

  • Method

    The paper proposes softmax-free self-attention whose linear variant has Petrov-Galerkin-comparable approximation capacity, together with Galerkin projection-type layer normalization.

  • Results

    The attention-based operator learner surpasses Fourier Neural Operator evaluation accuracy and recovers coefficients from extremely noisy measurements that traditional methods or FNO cannot accomplish.

  • Takeaways & Limitations

    The method provides a simple attention-based operator learner for both PDE solution operators and inverse coefficient identification operators.

  • Takeaways & Limitations

    The approach assumes operators have certain low-dimensional attributes, and full-resolution attention is inefficient for two-dimensional problems.

Abstract

from arXiv · show

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.

1 Introduction

The paper adapts Transformer attention to PDE operator learning, removing softmax and introducing a Galerkin-inspired normalization to improve efficiency, stability, and accuracy. Experiments cover three benchmark problems, including inverse coefficient recovery from noisy measurements.

  • Motivation: Operator learners map between function spaces, enabling evaluation of many PDE instances without retraining or collocation points.This contrasts with function learners that approximate a single instance under fixed conditions.
  • Approach: The paper applies Transformer self-attention to data-driven operator learning and modifies the attention mechanism for this setting.The attention mechanism is used to capture long-range interactions in discretized input and output functions.
  • Approach: Removing softmax permits Galerkin-type attention to represent a Petrov-Galerkin projection, with approximation capacity independent of sequence length.The argument uses Hilbert-space interpretations and a discrete LBB condition.
  • Approach: A Galerkin projection-type layer normalization addresses instability in softmax-free models while allowing learnable scaling to propagate through encoder layers.The normalization is presented as a cheap diagonal alternative to the normalizations derived in the Petrov-Galerkin proof.
  • Experiments: The proposed operator learner recovers coefficients from noisy measurements that traditional methods or FNO cannot accomplish.The paper combines the new attention operators with FNO for PDE solution-operator benchmarks.
  • Experiments: Three benchmark problems show improved computational and memory efficiency and accuracy over attention models using conventional softmax normalization.The experiments include PDE solution operator learning and inverse coefficient identification.

2 Related Works

Related work spans operator learners for PDEs, attention variants, and Transformer approximations. Prior approaches include graph and convolutional models, DeepONet, additive attention, kernel interpretations, and Nyström approximations.

  • Operator learners for PDEs: PDE operator-learning research includes graph neural networks, DeepONet, multilevel-grid kernel methods, FNO, dimension reduction, and fixed-discretization convolutional networks.These approaches target solution maps or operators with differing architectural and discretization strategies.
  • Attention mechanism and variants: Earlier additive attention suffered from vanishing gradients, while an early softmax-free variant retained a sigmoid and performed worse than its softmax counterpart.The cited comparison motivates examining alternative attention formulations.
  • Various studies on Transformers: Transformer-related work motivates Galerkin reformulations, Fourier-type attention without softmax, Nyström kernel approximations, and alternative normalization schemes.These connections provide theoretical and architectural precedents for the paper’s attention adaptations.

3 Operator learning related to PDEs

The paper formulates PDE operator learning as approximating mappings between function spaces, including forward solution operators, inverse coefficient maps, and nonlinear initial-value evolution. Training uses sampled input-response pairs on discretized grids, with mesh-independent approximation as the objective.

  • Operator formulation: The target is a densely defined operator T: H1 → H2 between Hilbert spaces over a bounded spatial domain.The operator is associated with physical problems formulated through PDEs.
  • Parametric PDE: A parametric PDE maps a varying coefficient a to its solution u through La(u) = f, defining a forward operator T: A → X.The right-hand side is fixed in this formulation.
  • Inverse problem: The inverse coefficient-identification operator maps noisy steady-state measurements ũ back to the coefficient a.This reverses the forward solution mapping and is treated as a separate operator-learning problem.
  • Nonlinear initial-value problem: A nonlinear initial-value problem defines direct inference from u0 to a later solution state u(t1), with t1 much larger than the traditional integrator step size.The operator maps H0 to H.
  • Training and discretization: The model is trained from N sampled input-response pairs, and its approximation is evaluated on discrete-grid observations.The supplied formulation specifies sampled data and practical discrete-grid evaluation.
  • Training and discretization: The learned approximation should remain accurate independently of mesh size, while sampled coefficients represent a compactly supported input distribution.The loss may include a problem-dependent regularizer, such as an energy-based term in Darcy flow.

4 Attention-based operator learner

The operator learner uses simple attention variants that replace softmax with mesh-weighted normalization and interpret attention through integral transforms or Petrov–Galerkin projections. Fourier-type attention has quadratic complexity in sequence length, while Galerkin-type attention admits a linear-complexity formulation and a learnable projection interpretation.

  • Operator learner architecture: The encoder uses shared feature extractors, recurrent positional coordinates, and problem-dependent decoders to map discretized inputs to target representations.Feedforward networks serve one-dimensional grids, interpolation-based CNNs serve two-dimensional domains, and decoders vary with target regularity.
  • Simple self-attention: Simple attention removes softmax, applies mesh-weighted normalization, and permits scaling to propagate through encoder layers.The formulation uses trainable projections Q, K, and V, with head-wise normalization applied before dot products.
  • Fourier-type attention: Fourier-type attention computes a mesh-weighted integral transform with a learnable nonsymmetric kernel, linking scaled dot-product attention to Fourier-type kernel operators.Its output approximates integration of value features against a kernel formed from query and key feature maps.
  • Computational complexity: Both matmuls in Fourier-type attention have complexity O(n^2d), whereas Galerkin-type attention reorganizes the products with complexity O(nd^2).The Galerkin form computes Q(K^⊤V)/n rather than forming the sequence-by-sequence product first.
  • Galerkin-type attention: Galerkin-type attention recasts forward propagation as a learnable Petrov–Galerkin projection, with Q as trial space, K as test space, and V representing basis functions.The method can also be viewed as dynamically updating the basis through responses between value and key features.
  • Approximation and scalability: The approximation estimate scales with sequence length when the lower bound in the stability condition is independent of n.This condition is associated with the discrete LBB or BNB condition in Galerkin methods.

5 Experiments

The experiments compare Fourier and Galerkin Transformers with softmax-based attention and FNO baselines across three PDE-related operator-learning tasks. The proposed attention models reduce computational costs and improve evaluation accuracy, with normalization effects depending on data scaling.

  • Experimental setup: The study compares Fourier and Galerkin Transformers against softmax-normalized attention variants and Fourier Neural Operator baselines across three PDE-related tasks.The models use matched parameter quotas and common training protocols, with additional baseline results taken from prior work.
  • Experimental setup: A modified projection-matrix initialization improves evaluation performance by up to 50% and improves training stability in certain operator-learning tasks.The modification is described as a small change to Xavier initialization intended to address the lack of softmax normalization.
  • Efficiency: Fourier-type attention reduces memory usage by 40%–50% versus softmax attention, while Galerkin attention provides up to a 120% speed boost over standard linear attention in certain tests.The comparison uses memory, FLOP, complexity, and training iteration speed measurements.
  • 5.1 Example 1: viscous Burgers’ equation: For viscous Burgers’ equation, attention-based learners achieve resolution-invariant performance similar to FNO1d, and the Galerkin projection-type normalization substantially outperforms regular layer normalization on unnormalized data.The benchmark uses a periodic boundary condition and Gaussian Random Field-sampled initial conditions.
  • 5.2 Example 2: Darcy flow: For the Darcy interface problem, attention-based learners achieve average evaluation results 30% to 50% better than the fine-grid FNO2d baseline, while normalization choice has little effect on normalized data.The learned operator maps random-interface coefficients to weak solutions, with a finite-dimensional treatment of the function-space embedding.
  • 5.3 Example 3: inverse coefficient identification for Darcy flow: In inverse Darcy coefficient identification with noisy inputs, FNO2d struggles with high-frequency interfaces, whereas the attention-based learner can recover structurally difficult inverse responses.The inputs include added noise, and the inverse problem is described as ill-posed.

6 Conclusion

The paper proposes a general simple-attention operator learner for PDE solution and inverse coefficient operators. It reports higher benchmark evaluation accuracy than FNO while identifying scope and computational limitations.

  • 6 Conclusion: The proposed attention-based operator learner approximates both PDE solution operators and inverse coefficient-identification operators.The conclusion describes the network as versatile across these two operator-learning settings.
  • 6 Conclusion: Benchmark evaluation accuracy surpasses that of the current state-of-the-art Fourier Neural Operator.This conclusion is stated for the benchmark problems considered in the paper.
  • 6 Conclusion: The approach assumes low-dimensional attributes in the learned operator responses, is inefficient at full resolution for 2D attention, and is non-causal in its linear variant.These limitations constrain applications involving nonsmooth L∞ subsets and decoder or autoregressive settings.

7 Broader Impact

The paper connects Galerkin-type attention to broader directions in numerical analysis and scientific applications. It also presents theoretical guidance for feature maps that preserve physical structure.

  • The method approximates an inverse coefficient identification problem with extremely noisy measurements, which traditional nonlinear iterative methods could not attain.
  • The attention interpretation may support applications in geoscience and medical imaging.
  • The proof guides feature-map design toward structures such as differentiability and physical invariance.

Appendices of Choose a Transformer: Fourier or Galerkin

Table 4 lists the paper’s notations in approximate chronological order and gives their meanings.

  • Table 4 organizes the notations used throughout the work and records their meanings.

B Network structures

The network structures combine attention-based encoder layers with resolution-changing convolutional components and problem-dependent decoders. Inputs are downsampled to coarse grids before attention and later reconstructed toward the target resolution.

  • Network structures: The forward Darcy model uses the attention-based operator learner shown in Figure 5, while Figure 4 corresponds to Example 5.1.
  • Downsampling CNN: Downsampling CiNNs reduce fine-grid inputs to coarse-grid latent functions and expand channels to the attention layers’ hidden dimension.
  • Downsampling CNN: Bilinear interpolation replaces pooling in the downsampling CiNN, whose convolution block removes batch normalization.
  • Attention encoder: Coarse-grid positional encodings are concatenated with latent representations before flattening and scaled dot-product attention.
  • Upsampling CNN: The upsampling CiNN reshapes encoder outputs, performs two interpolations with an intermediate convolution, and concatenates fine-grid positional encodings before decoding.

C.1 Training and evaluation setup

The experiments use fixed training protocols, controlled randomness, and reported computational costs across three operator-learning examples. Training uses 1cycle scheduling with ADAM, while evaluation includes seed-based error bands.

  • Training setup: Models are trained for 100 epochs with mini-batch ADAM, batch size 8, and a 1cycle learning-rate schedule.The schedule warms up to lrmax at epoch 30; lrmax is 10^-3 except 5 × 10^-4 for ST and FT in 2D.
  • Computational cost: Attention-based operator learners require only a fraction of the training time used by MLP-based operator learners.
  • Datasets: Each of the three examples contains 1024 training samples and 100 testing samples without repeated functions across the two sets.
  • Regularization: An H1-seminorm regularization term is used when target smoothness is known a priori, with γ set separately for Examples 5.1 and 5.2.The reported choices are γ = 0.1h for Example 5.1 and γ = 0.5h for Example 5.2.
  • Reproducibility: The experiments use a fixed random seed and deterministic cuDNN, while error bands are computed from 10 different seeds.

C.2 Experimental details for Example 5.1

Example 5.1 learns a discretized viscous Burgers’ solution operator from Gaussian-random initial data across three grid resolutions. The experiments examine normalization and initialization choices, while related Darcy-flow details and evaluation results are also reported.

  • Data preparation: The viscous Burgers’ equation is solved spectrally at high resolution, then downsampled to resolutions n = 512, 2048, and 8192 for operator learning.The viscosity is ν = (2π)^−10.1, and the finest-grid solution uses n = 8192.
  • Discrete operator: The learned discrete operator maps interpolated initial data in X_h to the downsampled solution at final time.The operator is written as T_h: X_h → X_h, with interpolation I_h and finest-grid downsampling Π_h.
  • Normalization and stability: The new Galerkin projection-type normalization reaches the common ≈1 × 10^−3 evaluation-error range, whereas regular normalization fails to converge under the tested setup.With regular normalization, the metric peaks around 2 × 10^−2 after warmup; diagonal initialization supports convergence for the new rule.
  • Normalization and stability: The Galerkin projection-type normalization is motivated by preserving the deterministic energy decay of a Burgers’ solution instance.The authors contrast this scale-preserving behavior with positionwise normalization and report that instance normalization deteriorates training stability.
  • Example 5.2: For interface Darcy flow, the coefficient takes values 12 and 3 almost surely, with a random interface pattern and fixed forcing f ≡ 1.The discrete forward problem uses a second-order five-point finite-difference scheme on a 421 × 421 grid.
  • Example 5.2: The Darcy-flow operator maps diffusion coefficients to unique weak solutions, with pointwise Gaussian normalization applied to inputs and targets.The network includes a non-trainable normalizer transforming values at each grid point toward N(0, 1).
  • Evaluation: A randomly selected Darcy-flow test sample has relative L2-error 6.454 × 10^−3, while four Burgers’ test samples have average relative error 1.079 × 10^−3.The Darcy figure also displays the input coefficient and the L∞-error distribution.
  • Limitations: The method remains less accurate than traditional discretization methods optimized for a single instance, with larger errors near sharp coefficient contrasts.Adaptive allocation of degrees of freedom based on a posteriori local error is identified as future work.

C.4 Experimental details for Example 5.3

Example 5.3 studies inverse recovery of a coefficient from noisy fine-grid solution measurements, a problem described as less stable than the forward map. The Galerkin Transformer is analyzed through approximation theory and evaluated under varying noise and resolution constraints.

  • Inverse problem: The inverse task reconstructs a coarse-grid coefficient from noisy solution measurements sampled on a fine grid.The noise level ε controls perturbations in the measurements, and ε = 0.1 corresponds to 10% noise.
  • Inverse problem: Noise makes the learned inverse operator not well-posed, and inverse problems are generally less stable than solving the corresponding forward problem.The forward problem maps coefficients to solutions, whereas the inverse map recovers coefficients from measurements.
  • Evaluation: The reported Galerkin Transformer reconstruction errors are 0.0160, 0.0292, and 0.0885 for 0%, 1%, and 10% noise, respectively.Inputs are shown on a 211 × 211 grid and recovered coefficients on a 71 × 71 grid.
  • Resolution limitation: The model reconstructs only coarse-grid coefficients because upsampling interpolation loses approximation order when high-frequency target components dominate.The architecture resolves low-frequency trends well but is limited by the coarse grid for low-regularity targets such as L∞ functions.
  • Noise limitation: Evaluation is unstable under noise when the model is trained on clean data, especially for the inverse problem’s large Lipschitz constant amplified by the black-box model.The authors recommend adding noise during training for inverse coefficient identification.
  • Approximation theory: Theorem 4.3 shows that the proposed linear attention variant can explicitly replicate a Petrov–Galerkin projection in the current latent representation subspace.The proof is formulated under a Hilbertian setup and connects sequence-length-uniform bounds to operator approximation.
  • Approximation theory: The attention interpretation uses different, learnable test and trial spaces, bringing Petrov–Galerkin projection into the model.The Galerkin-type normalization acts as a cheap learnable alternative to Gram-matrix inverses in the explicit projection formula.
  • Dynamic basis generation: Positional encodings dynamically update the approximation subspaces; without them, attention remains a combination of the current learned basis or subspace.The identity operator can still be learned exactly in the approximation subspace under stated caveats.

D.3 Proof of Theorem 4.3

The proof connects Galerkin-type attention to a Petrov-Galerkin projection and shows that its approximation capacity can match the best approximation in the value space. The argument also motivates layer-wise basis updates and Galerkin-type normalization.

  • Setup: The proof represents query, key, and value features as bases of finite-dimensional approximation spaces with distinct nodal and degree-of-freedom representations.The value space Qh is generated by columns of Q, while the key space Vh is generated by columns of V.
  • Setup: The learnable map gθ combines Galerkin-type attention with a pointwise universal approximator to map latent representations into the value space Qh.Its parameters include updated projections for query, key, and value features.
  • Approximation argument: Theorem 4.3 bounds the attention learner through the best approximator in Qh and a Petrov-Galerkin projection problem.The best approximator fh depends on the value space, whereas the learnable map must represent the corresponding projection.
  • Approximation argument: Riesz representation and finite-dimensional saddle-point arguments produce a value-to-key map and a solvable block system for the Petrov-Galerkin projection.The proof establishes existence through the bilinear form, matrix invertibility, and positive definiteness.
  • Representation: The learned attention can represent the projection after permuting the basis matrix and updating the query, key, and value parameters.This establishes the learner’s capacity to realize the best approximator used in the theorem.
  • Normalization and scope: The layer normalization scheme is motivated as a diagonal approximation to the projection-related normalizations, while globally supported learned bases make the computational spaces effectively mesh-free.The mesh-free characterization is noted as relevant to future study of zero-shot super-resolution.

D.5 Auxiliary results

The auxiliary results formalize the finite-dimensional Hilbert-space setting, derive the saddle-point and matrix formulations, and analyze how sequence length affects stability constants. They identify softmax normalization as a source of sequence-length dependence.

  • Finite-dimensional setting: The latent value space Vh is embedded in a nodal representation space Yh, with basis functions evaluated at discretization points and an induced continuous inner product.The construction uses d linearly independent basis functions with d < n grid points.
  • Solvability: The matrix BGB⊤ is invertible when B has full row rank and G is symmetric positive definite, supporting solvability of the auxiliary system.The proof establishes this through a Cholesky factorization and rank preservation.
  • Variational formulation: The saddle-point formulation is equivalent to an operator equation system coupling an inner-product equation with a bilinear-form constraint.The auxiliary lemma then translates this system into a block matrix formulation using basis coefficients.
  • Stability: The lower-bound argument uses the singular values of the bilinear-form matrix B together with norm-equivalence constants between functions and coefficient vectors.The assumptions include norm equivalence for the relevant finite-dimensional spaces.
  • Sequence-length dependence: For globally supported orthonormal bases, the norm-equivalence constants are approximately 1, while locally supported bases introduce mesh-dependent scaling that is offset by the singular-value scaling.The resulting final bound is sequence-length-independent under the stated scaling behavior.
  • Sequence-length dependence: Softmax applied along the sequence dimension makes the norm-equivalence constant cV depend on n, yielding an estimate with an unavoidable sequence-length-related constant.The analysis uses a one-dimensional L2 setting and an additional weak-derivative assumption.
Loading 2105.14995v4…