[Paper Review] Choose a Transformer: Fourier or Galerkin

Shuhao Cao|arXiv (Cornell University)|May 31, 2021

Model Reduction and Neural Networks97 references58 citations

TL;DR

The paper introduces Fourier-type and Galerkin-type self-attention variants (without softmax) for operator learning in PDEs, proving their approximation capacity and demonstrating efficiency and accuracy gains over softmax-based transformers on PDE benchmarks.

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.

Motivation & Objective

• Motivate data-driven operator learning for PDEs and the need for mesh-invariant, resolution-invariant operators.
• Introduce a softmax-free self-attention variant that aligns with Hilbert-space and Petrov–Galerkin principles.
• Develop the Galerkin Transformer and a corresponding layer normalization to stabilize training on unnormalized data.
• Demonstrate improved training efficiency and accuracy on PDE benchmarks compared to softmax-based attention.

Proposed method

• Define simple self-attention variants without softmax and with mesh-length weighted normalization (Galerkin-type and Fourier-type).
• Interpret attention as matrix-free approximations to integral transforms or Petrov–Galerkin projections in Hilbert spaces.
• Show that Galerkin-type attention corresponds to a learnable Petrov–Galerkin projection with n-independent approximation; provide Céa-type error bound (Theorem 4.3).
• Propose a layer normalization scheme mimicking Petrov–Galerkin projection to stabilize training and propagate scaling across layers.
• Enrich latent bases dynamically via feed-forward networks and positional encodings to adapt to operator responses."
• Compare FT and GT against softmax-based ST and LT variants on PDE tasks, highlighting memory and speed benefits.

Experimental results

Research questions

• RQ1 Can softmax-free attention achieve comparable approximation power to Petrov–Galerkin projections in operator learning for PDEs?
• RQ2 Does Galerkin-type attention enable stable, scalable training on unnormalized data and offer efficiency gains over traditional softmax attention?
• RQ3 How do Fourier-type and Galerkin-type attentions perform on benchmark PDE operator learning tasks (viscous Burgers’, Darcy flow, inverse coefficient identification)?
• RQ4 What role do layer normalization and dynamic basis updates play in improving accuracy and stability of operator learners?

Key findings

• Softmax-free Galerkin-type attention can explicitly represent a Petrov–Galerkin projection and its approximation capacity is independent of sequence length (Theorem 4.3).
• A new Galerkin projection-type layer normalization enables stable scaling through attention layers and improves physical property tracking (e.g., energy decay).
• The Galerkin Transformer and Fourier Transformer reduce training memory and can speed up training relative to softmax-based counterparts in operator learning tasks.
• Experiments on viscid Burgers’ equation, interface Darcy flow, and inverse interface coefficient identification show improved accuracy and efficiency with GT/FT.
• Memory usage is reduced by 40-50% for Fourier-type attention versus softmax, and Galerkin-type attention can offer significant speedups over linear attention in some cases.
• The authors provide open-source PyTorch code for reproducing results.