PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformer for Partial Differential Equations' Operator Learning

Zijie Li, Kazem Meidani, Amir Barati Farimani

arXiv:2205.13671v3cs.LGcs.AI

TL;DR

High-resolution numerical solutions for complex spatiotemporal PDEs are computationally expensive, while learned operators can be restricted to training resolutions, geometries, or grids. OFormer uses attention-based operator learning and shows competitive benchmark performance with flexibility across discretizations and irregular sampling.

  • Problem

    High-resolution discretization makes complex spatiotemporal PDEs computationally expensive, while learned operators may remain restricted to the resolutions, geometries, and grids seen during training.

  • Method

    OFormer is a point-based attention architecture combining cross-attention, self-attention, and point-wise MLPs to query outputs across varying input and output discretizations.

  • Results

    OFormer is competitive on benchmark PDE operator-learning problems, outperforms a GNN without mesh augmentation on irregular-grid Poisson problems, and predicts reasonably accurate full-resolution outputs from 25% randomly sampled inputs.

  • Takeaways & Limitations

    The framework supports varying numbers of input and query points, uniform and non-uniform grids, and generalization to diverse PDE instances without inputting system parameters.

  • Takeaways & Limitations

    Learned function mappings are restricted to the specific resolution and geometry observed during training.

Abstract

from arXiv · show

Data-driven learning of partial differential equations' solution operators has recently emerged as a promising paradigm for approximating the underlying solutions. The solution operators are usually parameterized by deep learning models that are built upon problem-specific inductive biases. An example is a convolutional or a graph neural network that exploits the local grid structure where functions' values are sampled. The attention mechanism, on the other hand, provides a flexible way to implicitly exploit the patterns within inputs, and furthermore, relationship between arbitrary query locations and inputs. In this work, we present an attention-based framework for data-driven operator learning, which we term Operator Transformer (OFormer). Our framework is built upon self-attention, cross-attention, and a set of point-wise multilayer perceptrons (MLPs), and thus it makes few assumptions on the sampling pattern of the input function or query locations. We show that the proposed framework is competitive on standard benchmark problems and can flexibly be adapted to randomly sampled input.

1 Introduction

PDE solution operators are difficult to approximate efficiently because conventional solvers require costly high-resolution discretization, while many learned models remain tied to training resolutions and geometries. OFormer addresses these constraints with attention-based operator learning that supports flexible input and query discretizations.

  • 1 Introduction: High-resolution discretization in time and space makes solving complex spatiotemporal PDEs computationally expensive.Numerical schemes discretize the solution space and solve a finite-dimensional problem.
  • 1 Introduction: Many data-driven PDE solvers impose architecture-specific inductive biases, such as convolutions for structured grids or graph layers for unstructured local relations.These methods learn solutions directly from observations without prior knowledge of the governing PDE.
  • 1 Introduction: Learned function mappings can be restricted to the specific resolution and geometry observed during training.This limits their applicability when discretization settings change.
  • 1 Introduction: Existing data-driven operators often require matching input and output discretizations, fixed input grids, or equi-spaced grid points.These restrictions arise despite promising performance on PDE-solving tasks.
  • 1 Introduction: OFormer uses cross-attention for discretization-invariant output queries and a Transformer architecture that handles varying numbers of points on uniform or non-uniform grids.Its contributions also include a latent time-marching scheme that reduces time-dependent PDE modeling to a latent ODE with a fixed time interval.

2 Related works

Related work spans deep-learning PDE solvers, neural operators, and attention-based models for physical systems. OFormer builds on these directions while using latent-space recurrence to support scalable temporal modeling.

  • Learned PDE Solvers: Physics-informed methods use PDE supervision to approximate solutions in aggregated space-time domains and are generally mesh-agnostic.The passage characterizes these methods as accurate while noting their PDE-based supervision.
  • Time-marching in the latent space: A common time-dependent strategy maps the solution at time t to the next time step through an Encoder-Process-Decoder scheme.Latent time-marching is introduced as an alternative architecture for temporal modeling.
  • Time-marching in the latent space: The proposed latent time-marching architecture propagates dynamics in latent space, allowing future-time unrolling during training while improving attention-encoder scalability.Latent propagation avoids repeatedly encoding the observable state.
  • Neural Operators: Neural operator research includes DeepONet’s operator-approximation framework and architectures based on iterative learnable integral kernels.These approaches approximate mappings between function spaces rather than only fixed finite-dimensional outputs.
  • Transformer for physical system: Attention has been explored for physical-system modeling, with prior work dividing broadly into orthogonal lines of using attention to capture structures or patterns.This places OFormer within a broader effort to apply Transformer-style mechanisms beyond natural language processing.

3 Method

OFormer combines attention-based encoders with cross-attention and recurrent MLP propagation to learn PDE solution operators across flexible input and query discretizations. Its design uses linearized attention for efficiency, query-location encoding for arbitrary outputs, and latent-space time marching for scalable trajectory training.

  • Attention mechanism: Standard softmax attention computes QK^T explicitly, giving O(n^2d) complexity for sequences of length n.The quadratic interaction is prohibitively expensive for long sequences.
  • Attention mechanism: Linear attention interprets matrix columns as learned basis-function samples and computes Z = Q(K^T V)/n without explicitly forming QK^T.Associative matrix multiplication enables efficient Fourier- and Galerkin-type attention; instance normalization is used in the model.
  • Attention mechanism: Cross-attention encodes queries from output coordinates, decoupling arbitrary query locations from the input sampling grid.The input-side keys and values are defined on D(x), while query basis functions are defined on D(y).
  • Attention mechanism: The cross-attention output is related to DeepONet, while appended point-wise FFNs improve performance and propagating MLPs are necessary for time-dependent PDEs.Attention-layer parameter count does not depend on input discretization, allowing varying grids without retraining.
  • Model architecture: The encoder processes sampled input values and coordinates, encodes query coordinates, and transfers system information to queries through cross-attention.Coordinate features use random Fourier projections to alleviate spectral bias in coordinate-based neural networks.
  • Model architecture: Latent-space recurrent propagation calls the encoder once, reducing memory usage and enabling scalable fully unrolled training for time-dependent PDEs.This approach propagates dynamics in latent space rather than the observable space.

4 Experiment

The experiments evaluate OFormer on uniform and irregular grids, representation learning, and architectural ablations. Results show competitive benchmark performance, flexibility across discretizations and sampling densities, generalization to unseen geometries, and meaningful latent representations, while attention choices remain problem-dependent.

  • Experimental setup: The experiments cover standard uniform-grid PDE benchmarks, irregular-grid problems, latent representation analysis, and ablations of architectural choices.The benchmark suite includes Burgers’, Darcy flow, Navier-Stokes, Poisson, and compressible flow around an airfoil.
  • Performance on equi-spaced grid: OFormer is generally not superior in small-data or smooth-target settings, where spectral inductive biases can be advantageous.For complex Navier-Stokes instances in the data-sufficient regime, it performs on par with MWT on NS2-full using substantially fewer parameters and similarly across resolutions.
  • Application to irregular grids: On irregular grids, OFormer outperforms the GNN baseline without mesh augmentation and is more robust to changes in local mesh topology.Both models extrapolate to new geometries with reasonable accuracy under mesh augmentation, while OFormer also applies to highly irregular meshes.
  • Application to irregular grids: The airfoil experiment reports a predicted-momentum test root mean squared error of 15.469±0.074 and relative error of 8.17±0.06%.The model predicts the temporal evolution of compressible flow around the airfoil on highly irregular meshes.
  • Variable discretizations: OFormer handles variable numbers of input and query points with shared weights, predicting a full-resolution output from inputs containing 25% of grid points.The 25% sparse input corresponds to 150% of the maximum dropping ratio used during training.
  • Learned spatio-temporal representation: Latent representations encode Navier-Stokes viscosity information: PCA reveals color trends, and linear regression effectively predicts viscosity from latent vectors.The model learned these representations without being given coefficient values during prediction training.
  • Ablation study: Ablations find that attention type and normalization target are problem-specific, larger models help in data-sufficient regimes, and removing RoPE significantly worsens predictions.The larger OFormer improves the previous best NS-full result from 0.0667 to 0.0625; RoPE preserves important relative position information.

5 Conclusion

OFormer is a fully point-based attention architecture for PDE solution-operator learning that performs competitively while supporting flexible input/output discretizations. The authors also identify data requirements, computational cost, and long-term error as important limitations.

  • OFormer is a fully point-based attention architecture for learning PDE solution operators.
  • The model shows competitive performance on PDE operator-learning problems and flexibility across input/output discretizations.
  • OFormer can generalize to diverse PDE instances without system parameters and learn meaningful representations for system identification.
  • The model needs sufficient data for optimal performance because it makes few assumptions about grid structure.
  • Linear attention remains relatively computationally expensive on higher-resolution grids, while long-term transient predictions accumulate unneglectable error.

A Model implementation details

The implementation uses PyTorch and represents the model through linear layers, attention heads, and separate input and query encoder architectures. The listed architectures include self-attention, cross-attention, propagator, and decoder components.

  • All models are implemented in PyTorch, with linear layers denoted by their input and output dimensions.The notation [d1, d2] corresponds to nn.Linear(d1, d2).
  • Attention layers are described by head dimension multiplied by the number of attention heads.
  • The input encoder architecture includes self-attention and feed-forward-network components.The architecture table identifies SA as self-attention and FFN as feed-forward network.
  • The query encoder architecture uses cross-attention, while the propagator and decoder are specified separately.The architecture tables identify CA as cross-attention and provide separate propagator and decoder layouts.

A.2 Hyperparameter for irregular grid problem

The irregular-grid implementation uses point-based encoders and specialized training choices for different problems. These include node-type embeddings, curriculum training, initialization adjustments, and normalization choices.

  • For irregular-grid problems, the query encoder adds a self-attention layer after cross-attention.
  • The airfoil model adds learnable embeddings for boundary, open-area, and airfoil node types.
  • Training on 2D Navier-Stokes and airfoil uses curriculum learning that initially truncates the prediction horizon with γ ≈ 0.5.
  • For Burgers’ equation, orthogonal query/value projection initialization with gain 1/d and an added diagonal constant 1/d was used.The latent dimension at each attention head is d.
  • Instance normalization is generally used, but layer normalization is adopted for Electrostatics/Magnetostatics because instance normalization is unstable with 60–250 grid points.

B Baseline details

The benchmark compares OFormer with FNO, MWT, and Galerkin/Fourier Transformer implementations under specified computational settings. The authors caution that exact speedups over numerical solvers are difficult to determine, despite efficient learned time-dependent solvers.

  • The benchmarks use official implementations of FNO, MWT, and Galerkin/Fourier Transformer with minimal setting changes.
  • For 2D Navier-Stokes, FNO uses 12 Fourier modes and width 32, while G.T./F.T. uses hidden dimension 96 and extended training.G.T./F.T. training is extended to 200 epochs, or 150 epochs on the 10,000-sample dataset.
  • The efficiency benchmark uses 10,000 training samples, time horizon T = 30, batch size 10, and an RTX-3090 GPU.
  • Exact acceleration over traditional numerical solvers is difficult to conclude because GPU support, optimization, accuracy, and grid choices vary.
  • Learned solvers for time-dependent systems generally tolerate time steps 3–4 orders of magnitude larger than numerical solvers.

C Further ablation study

The ablations examine OFormer’s sensitivity to discretization, initialization, data quantity, and Random Fourier Features. Results show robustness benefits under irregular sampling, but performance depends on data and design choices, with temporal errors accumulating over time.

  • Sparse reconstruction and prediction: Randomly dropping input points during training usually worsens performance, but improves robustness to irregular discretization and permits reasonable prediction from 25% of input points.For NS-mix and NS-mix2, random dropping has little influence on performance.
  • Influence of initialization: Without normalization, projection-matrix initialization is highly sensitive on Burgers’ equation, and standard orthogonal or Xavier initialization can cause performance to deteriorate.The reported scaling uses σ = 1/d and adds diagonal bias δI to the original weight matrix.
  • Influence of data: OFormer performance improves almost linearly as the data size grows exponentially, with especially significant gains in the 500–5000 sample regime.The dependence is particularly pronounced for complicated problems such as Navier–Stokes.
  • Influence of Random Fourier Feature: Random Fourier Features boost performance on all three studied problems by broadening the frequency spectrum of learned bases.Without RFF, the learned bases exhibit relatively smoother patterns than with RFF.
  • Temporal trend of error: Prediction error generally accumulates exponentially with time, while compressible-flow error shows some periodicity; more training data reduces error without changing this trend.The authors identify mitigating compounding prediction error as an important future direction.

D.1 Data generated on equi-spaced mesh

The equi-spaced-mesh benchmarks use standard PDE datasets generated on regular simulation grids and downsampled to different resolutions. They cover Burgers’ equation, Darcy flow, and Navier–Stokes operator-learning tasks.

  • Benchmark construction: The benchmark datasets are generated on equi-spaced simulation grids and downsampled to create different resolutions.These datasets are used as benchmarks in prior operator-learning studies.
  • Burgers’ Equation: For Burgers’ equation, the learned operator maps an initial condition u0(x) to the solution u(·, 1) under periodic boundary conditions.The viscosity is 0.1, and data are generated with a split-step method on an 8192-resolution grid.
  • Darcy Flow: For Darcy flow, the learned operator maps the diffusion coefficient a(·) to the steady-state solution u(·).The data use zero Neumann boundary conditions and are generated with a second-order finite-difference solver on a 421 × 421 grid.
  • Navier-Stokes Equation: For 2D Navier–Stokes, the learned operator maps vorticity over an initial time interval to future vorticity states.Lower viscosity coefficients produce more chaotic dynamics.

D.2 Data generated on non-equi-spaced mesh

The non-equi-spaced-mesh experiments include electrostatic and magnetostatic Poisson problems and compressible flow around an airfoil. Their operators map physical inputs to solution fields over nonuniform settings.

  • Poisson equation: For Poisson problems, the operator maps {ρ, ϵ} to electric potential U or {Iz, µ} to magnetic potential Az under Dirichlet boundary conditions.The permittivity ϵ and permeability µ are held constant in the studied problem.
  • Poisson equation: Poisson-equation data are generated with a finite-element method using quadratic triangle elements implemented in FEniCS.The experiments use pre-generated data from Lötzsch et al. (2022).
  • Euler equation of compressible flow: For compressible flow, the operator maps input fields {u, ρ, p} over an initial time interval to later velocity, density, and pressure fields.The problem imposes a no-penetration condition at the airfoil.
  • Euler equation of compressible flow: Unlike the referenced learned solver, this formulation is non-Markovian and uses a much larger time step of Δt = 0.192 instead of Δt = 0.008.The comparison contrasts direct mapping across time intervals with a Markovian one-step rollout.

E Further results and visualization

Additional visualizations report point-wise absolute-error-based predictions for Burgers’ equation, Darcy flow, Navier–Stokes systems, compressible flow, and a vector-field problem. The figures show model outputs across increasingly challenging flow settings and interpolated nonuniform points.

  • Visualization setup: The visualization section evaluates predictions using point-wise absolute error plots across several PDE problems.The reported figures include both scalar and vector-valued outputs.
  • Prediction visualizations: Figure 14 compares 1D Burgers’ predictions with ground truth, using green dotted lines for truth and red lines for predictions.Samples are randomly selected from the test set.
  • Prediction visualizations: Figures 15–18 visualize predictions for Darcy flow and Navier–Stokes test sets spanning Reynolds numbers around 20, 200, and 2000.The Navier–Stokes cases include NS1, NS2-full, and NS3.
  • Prediction visualizations: The compressible-flow visualization reports a test root mean squared error of 15.477 for predicted momentum in the plotted region.Momentum is computed by multiplying velocity by density at each node.
  • Prediction visualizations: Figure 20 presents interpolated predictions from nonuniform grid points for potential and the x and y components of a vector field.The three displayed quantities are potential, vector-field x component, and vector-field y component.
Loading 2205.13671v3…